324 lines
10 KiB
Python
324 lines
10 KiB
Python
|
from typing import Mapping, MutableMapping, Sequence, Iterable, List, Set
|
||
|
import pyfuse3
|
||
|
|
||
|
from iota import Iota, ProposedTransaction, Address, TryteString, Tag
|
||
|
from iota.crypto.addresses import AddressGenerator
|
||
|
from iota.crypto.types import Seed
|
||
|
from iota.codecs import TrytesDecodeError
|
||
|
|
||
|
from Crypto.Cipher import AES
|
||
|
from Crypto.Util.Padding import pad, unpad
|
||
|
|
||
|
import math
|
||
|
from pprint import pprint
|
||
|
import hashlib
|
||
|
import sys
|
||
|
import random
|
||
|
import time
|
||
|
|
||
|
import msgpack
|
||
|
import asyncio
|
||
|
import copy
|
||
|
|
||
|
import gzip
|
||
|
import secrets
|
||
|
|
||
|
CHUNKSIZE = 2187 - 1
|
||
|
|
||
|
def log(txt):
|
||
|
print("[-] "+str(txt))
|
||
|
|
||
|
class Atom():
|
||
|
def __init__(self) -> None:
|
||
|
self.type = "null"
|
||
|
|
||
|
def load(self, data) -> None:
|
||
|
if data[0]==True:
|
||
|
self.type = "milestone"
|
||
|
self.milestone = data[1:]
|
||
|
else:
|
||
|
if data[1]==True:
|
||
|
self.type = "dir"
|
||
|
else:
|
||
|
self.type = "file"
|
||
|
self.size = data[5]
|
||
|
self.name = data[2]
|
||
|
self.token = data[3]
|
||
|
self.milestoneIndex = data[4]
|
||
|
|
||
|
def dump(self) -> bytes:
|
||
|
# TODO: Delimiter?
|
||
|
# TODO: compression?
|
||
|
return msgpack.dump(self._dumpAsArray())
|
||
|
|
||
|
def _dumpAsArray(self):
|
||
|
if self.type=="milestone":
|
||
|
return [True]+self.milestone
|
||
|
else:
|
||
|
return [False, self.type, self.name, self.token, self.milestoneIndex, self.size]
|
||
|
|
||
|
class BlobChunk():
|
||
|
def __init__(self, data: bytes = b'', sealed: bool = False) -> None:
|
||
|
self.data = data
|
||
|
self.sealded = sealed
|
||
|
|
||
|
def getData(self) -> bytes:
|
||
|
return self.data
|
||
|
|
||
|
def append(self, data: bytes) -> None:
|
||
|
if len(data)+len(self.data) > CHUNKSIZE:
|
||
|
raise Exception("That to big!")
|
||
|
self.data += data
|
||
|
|
||
|
def getBytesLeft(self) -> int:
|
||
|
if self.sealed:
|
||
|
return 0
|
||
|
return CHUNKSIZE - len(self.data)
|
||
|
|
||
|
def seal(self) -> None:
|
||
|
self.sealed = True
|
||
|
|
||
|
def isSealed(self) -> bool:
|
||
|
return self.sealed
|
||
|
|
||
|
class TangleBlob():
|
||
|
def __init__(self, token: bytes, iotaApi: Iota) -> None:
|
||
|
self.token = token
|
||
|
self.iotaApi = iotaApi
|
||
|
self.preChunks = 0
|
||
|
self.chunks = []
|
||
|
m = hashlib.sha3_512()
|
||
|
m.update(self.token)
|
||
|
self.adressGen = AddressGenerator(Seed(m.digest()))
|
||
|
self.fetched = False
|
||
|
self.pushedNum = 0
|
||
|
|
||
|
def _requireFetched(self):
|
||
|
if not self.fetched:
|
||
|
self.fetch()
|
||
|
|
||
|
def _getKey(self, chunkNum: int) -> bytes:
|
||
|
m = hashlib.sha3_384()
|
||
|
m.update(self.token)
|
||
|
m.update(chunkNum.to_bytes(8, "little")) # 64 bits should be enought...
|
||
|
m.update(self.token)
|
||
|
return m.digest()
|
||
|
|
||
|
def _genBundle(self, data, addr) -> str:
|
||
|
txMsg = TryteString.from_bytes(data)
|
||
|
trans = ProposedTransaction(
|
||
|
address = addr,
|
||
|
value = 0,
|
||
|
tag = Tag("IOTAFS"),
|
||
|
message = txMsg
|
||
|
)
|
||
|
return self.iotaApi.prepare_transfer(
|
||
|
transfers = [trans],
|
||
|
inputs = [addr]
|
||
|
)['trytes']
|
||
|
|
||
|
def _dumpChunk(self, chunkNum: int, security_level=2) -> str:
|
||
|
key = self._getKey(chunkNum + self.preChunks)
|
||
|
data = self.chunks[chunkNum].getData()
|
||
|
cipher = AES.new(key[16:][:16], AES.MODE_CBC, key[:16])
|
||
|
ct_bytes = cipher.encrypt(pad(data, AES.block_size))
|
||
|
addr = self.adressGen.get_addresses(start=chunkNum + self.preChunks, count=1, security_level=security_level)['addresses'][0]
|
||
|
return self._genBundle(ct_bytes, addr)
|
||
|
|
||
|
def append(self, data: bytes, newBlock: bool = False) -> None:
|
||
|
self._requireFetched()
|
||
|
if len(self.chunks) and not newBlock:
|
||
|
bytesLeft = self.chunks[-1].getBytesLeft()
|
||
|
if bytesLeft:
|
||
|
leftChunk = data[:bytesLeft]
|
||
|
data = data[bytesLeft:]
|
||
|
self.chunks[-1].append(leftChunk)
|
||
|
while len(data):
|
||
|
chunk = data[:CHUNKSIZE]
|
||
|
self.chunks.append(BlobChunk(chunk))
|
||
|
data = data[CHUNKSIZE:]
|
||
|
|
||
|
def getChunkLen(self) -> int:
|
||
|
return self.preChunks + len(self.chunks)
|
||
|
|
||
|
def getSize(self) -> int:
|
||
|
if len(self.chunks):
|
||
|
return self.getChunkLen()*CHUNKSIZE - self.chunks[-1].getBytesLeft()
|
||
|
return self.preChunks
|
||
|
|
||
|
def read(self) -> bytes:
|
||
|
data = b''
|
||
|
for chunk in self.chunks:
|
||
|
data += chunk.getData()
|
||
|
return data
|
||
|
|
||
|
def _dump(self) -> str:
|
||
|
self.chunks[-1].seal()
|
||
|
data = ""
|
||
|
for c in range(len(self.chunks)-self.pushedNum):
|
||
|
num = c + self.pushedNum
|
||
|
data += self._dumpChunk(num, 2)
|
||
|
return data
|
||
|
|
||
|
def fetch(self, security_level=2, skipChunks: int = 0) -> None:
|
||
|
chunkNum = self.getChunkLen() + skipChunks
|
||
|
while True:
|
||
|
key = self._getKey(chunkNum)
|
||
|
cipher = AES.new(key[16:][:16], AES.MODE_CBC, key[:16])
|
||
|
addr = self.adressGen.get_addresses(start=chunkNum, count=1, security_level=security_level)['addresses'][0]
|
||
|
txHash = self.api.find_transactions(tags=[Tag("IOTAFS")], addresses=[addr])["hashes"]
|
||
|
if len(txHash)==0:
|
||
|
break
|
||
|
bundles = self.api.get_bundles(txHash[0])["bundles"]
|
||
|
for bundle in bundles:
|
||
|
for tx in bundle.transactions:
|
||
|
# TODO: Can we just strip the 9s and call it a day?
|
||
|
tryteStr = TryteString(str(tx.signature_message_fragment).rstrip("9"))
|
||
|
try:
|
||
|
ct_bytes = tryteStr.as_bytes()
|
||
|
except TrytesDecodeError:
|
||
|
ct_bytes = (tryteStr+"9").as_bytes()
|
||
|
self.chunks.append(BlobChunk(unpad(cipher.decrypt(ct_bytes), AES.block_size), True))
|
||
|
chunkNum += 1
|
||
|
self.pushedNum = len(self.chunks)
|
||
|
self.fetched = True
|
||
|
self._afterFetch()
|
||
|
|
||
|
def _afterFetch() -> None:
|
||
|
return
|
||
|
|
||
|
class TangleFileTreeElement(TangleBlob):
|
||
|
def __init__(self, token: bytes, iotaApi: Iota) -> None:
|
||
|
super(TangleFileTreeElement, self).__init__(token)
|
||
|
self.inodes = {}
|
||
|
self.atomStack = -1
|
||
|
|
||
|
def _afterFetch(self) -> None:
|
||
|
data = msgpack.load(self.read())
|
||
|
newAtoms = []
|
||
|
for i, elem in enumerate(reversed(data)):
|
||
|
atom = Atom()
|
||
|
atom.load(elem)
|
||
|
if atom.type == "milestone":
|
||
|
self._applyMilestone(atom.milestone)
|
||
|
break
|
||
|
newAtoms.append(atom)
|
||
|
self.atomStack = len(newAtoms)
|
||
|
for atom in reversed(newAtoms):
|
||
|
self._applyAtom(atom)
|
||
|
|
||
|
def _applyMilestone(self, milestone) -> None:
|
||
|
self.atomStack = 0
|
||
|
self.inodes = {}
|
||
|
for stone in milestone:
|
||
|
atom = Atom()
|
||
|
atom.load(stone)
|
||
|
self.inodes[atom.name] = atom
|
||
|
|
||
|
def _applyAtom(self, atom: Atom) -> None:
|
||
|
if atom.name in self.inode:
|
||
|
self.atomStack += 1
|
||
|
self.inodes[atom.name] = atom
|
||
|
if atom.type=="dir":
|
||
|
self.inodes[atom.name].elem = TangleFileTreeElement(atom.token, self.iotaApi)
|
||
|
elif atom.type=="file":
|
||
|
self.inodes[atom.name].elem = TangleFile(atom.token, self.iotaApi)
|
||
|
else:
|
||
|
raise Exception("How did such an atom get here?")
|
||
|
|
||
|
def _newAtom(self, atom: Atom) -> None:
|
||
|
self.append(atom.dump())
|
||
|
|
||
|
def getNameList(self) -> List(str):
|
||
|
self._requireFetched()
|
||
|
return list(self.inodes.keys())
|
||
|
|
||
|
def getInode(self, name: str) -> Atom:
|
||
|
self._requireFetched()
|
||
|
return self.inodes[name]
|
||
|
|
||
|
def mkdir(self, name: str) -> bool:
|
||
|
self._requireFetched()
|
||
|
if name in self.getNameList():
|
||
|
return False
|
||
|
atom = Atom()
|
||
|
atom.type="dir"
|
||
|
atom.name = name
|
||
|
atom.token = "GENTOKENHERE" # TODO: <
|
||
|
atom.milestoneIndex = 0
|
||
|
self._newAtom(atom)
|
||
|
self._applyAtom(atom)
|
||
|
return True
|
||
|
|
||
|
def mkfile(self, name: str) -> bool:
|
||
|
self._requireFetched()
|
||
|
if name in self.getNameList():
|
||
|
return False
|
||
|
atom = Atom()
|
||
|
atom.type="file"
|
||
|
atom.name = name
|
||
|
atom.token = "GENTOKENHERE" # TODO: <
|
||
|
atom.milestoneIndex = 0
|
||
|
self._newAtom(atom)
|
||
|
self._applyAtom(atom)
|
||
|
return True
|
||
|
|
||
|
def _updateFileSize(self, name: str, size: int):
|
||
|
self._requireFetched()
|
||
|
self.inodes[name].size = size
|
||
|
self._newAtom(self.inodes[name])
|
||
|
|
||
|
def _updateFileToken(self, name: str, token: bytes, size: int):
|
||
|
self._requireFetched()
|
||
|
self.inodes[name].token = token
|
||
|
self.inodes[name].size = size
|
||
|
self._newAtom(self.inodes[name])
|
||
|
|
||
|
def performMilestone(self):
|
||
|
stones = []
|
||
|
for atom in self.inodes:
|
||
|
stones.append(atom._dumpAsArray())
|
||
|
self.atomStack = 0
|
||
|
# TODO: Delimiter ?
|
||
|
# TODO: compression ?
|
||
|
data = msgpack.dump(stones)
|
||
|
self.append(data)
|
||
|
|
||
|
class TangleFile():
|
||
|
def __init__(self, name: str, parent: TangleFileTreeElement, iotaApi: Iota) -> None:
|
||
|
self.name = name
|
||
|
self.parent = parent
|
||
|
self.reflexiveAtom = parent.inodes[name]
|
||
|
self.size = self.reflexiveAtom.size
|
||
|
self.token = self.reflexiveAtom.token
|
||
|
super(TangleFile, self).__init__(self.token, iotaApi)
|
||
|
|
||
|
def write(self, offset: int, data: bytes):
|
||
|
if offset == self.size:
|
||
|
self.append(data)
|
||
|
else:
|
||
|
pass
|
||
|
|
||
|
|
||
|
class IotaFs():
|
||
|
def __init__(self, token) -> None:
|
||
|
self.api = Iota('https://nodes.thetangle.org:443', local_pow=True)
|
||
|
self.genesis = TangleFileTreeElement(token, self.api)
|
||
|
log("Fetching Genesis...")
|
||
|
self.genesis.fetch()
|
||
|
log("Retrieving reference to root")
|
||
|
if self.genesis.getNameList()!=["/"]:
|
||
|
if len(self.genesis.getNameList()):
|
||
|
# theres another directory in our genesis chain... WTF?!
|
||
|
raise Exception("Corrupted Genesis-Chain:"
|
||
|
+ "Unknown records for no root-directory in Genesis Chain: "+str(self.genesis.getNameList()))
|
||
|
else:
|
||
|
# we dont have a root yet, lets create one...
|
||
|
log("Unable to reference to root: Creating new root")
|
||
|
self.genesis.mkdir("/")
|
||
|
log("Successfully Mounted!")
|
||
|
|
||
|
|
||
|
def createNewFile(self, name) -> None:
|
||
|
pass
|