# iotaFS a.k.a iotaShitPoc

# Big TODOs:
# - map from binary file to position in chain
#   -> ability to read of size easily and read / write in chunks is 1000000% more efficient
# - remove old / big stuff from cache
# - async sync with tangle -> 1000% more efficiency
# - implement deltas for fileTree or switch to actuall tree with links

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

from errno import ENOENT
from fuse import FUSE, FuseOSError, Operations, LoggingMixIn
import stat
import os
import gzip
import secrets

MAX_CACHE_AGE = 60*60 # 1h

class IotaFS_BlobStore():
    def __init__(self, api=None):
        if api==None:
            self.api = Iota('https://nodes.thetangle.org:443', local_pow=True)
        else:
            self.api = api

    def _genBundles(self, data, addrIter, lenPerTx = 2187, txPerBundle = 1):
        msg = TryteString.from_bytes(data)
        bundles = []
        nextAddr = addrIter.__next__()
        for b in range(math.ceil(len(msg)/(lenPerTx*txPerBundle))):
            bundleMsg = msg[lenPerTx*txPerBundle*b:][:lenPerTx*txPerBundle]
            bundleTxs = []
            addr = nextAddr
            print("[addr] "+str(addr.with_valid_checksum()))
            nextAddr = addrIter.__next__()
            for t in range(math.ceil(len(bundleMsg)/lenPerTx)):
                txMsg = bundleMsg[lenPerTx*t:][:lenPerTx]
                bundleTxs.append(
                    ProposedTransaction(
                        address = addr,
                        value = 0,
                        tag = Tag("IOTAFS"),
                        message = txMsg
                    )
                )
            bundles.append(
                self.api.prepare_transfer(
                    transfers = bundleTxs,
                    inputs = [addr]
                )['trytes']
            )
        return bundles

    def _sendBundles(self, bundles):
        bundleRets = []
        for i,bundle in enumerate(bundles):
            print(str(int(i/len(bundles)*100))+"%")
            bundleRets.append(
                self.api.send_trytes(
                    trytes=bundle
                )
            )
        return bundleRets

    def uploadData(self, data, secret):
        print("Uploading...")
        m = hashlib.sha3_384()
        m.update(secret)
        m.update(data)
        sHash = m.digest()
        self.uploadDataRaw(data, sHash)
        return sHash

    def uploadDataRaw(self, data, sHash):
        trSeed = TryteString.from_bytes(sHash[16:])[:81]
        cipher = AES.new(sHash[:16], AES.MODE_CBC, sHash[22:][:16])
        ct_bytes = cipher.encrypt(pad(data, AES.block_size))
        addrIter = AddressGenerator(Seed(trSeed)).create_iterator(start = 0, step = 1)
        bundles = self._genBundles(ct_bytes, addrIter)
        self._sendBundles(bundles)

    def uploadTxt(self, txt, secret):
        data = str.encode(txt)
        return self.uploadData(data, secret)

    def getData(self, sHash):
        print("Downloading...")
        trSeed = TryteString.from_bytes(sHash[16:])[:81]
        cipher = AES.new(sHash[:16], AES.MODE_CBC, sHash[22:][:16])
        addrIter = AddressGenerator(trSeed).create_iterator(start=0, step=1)
        tryteMsg = ""
        for addr in addrIter:
            print("[addr] "+str(addr.with_valid_checksum()))
            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:
                    tryteMsg+=str(tx.signature_message_fragment)
        if tryteMsg == "":
            return b''
        tryteStr = TryteString(tryteMsg.rstrip("9"))
        try:
            ct_bytes = tryteStr.as_bytes()
        except TrytesDecodeError:
            ct_bytes = (tryteStr+"9").as_bytes()
        data = unpad(cipher.decrypt(ct_bytes), AES.block_size)
        return data

    def getTxt(self, sHash):
        return self.getData(sHash).decode("utf-8")

    def getSHash(self, data, secret):
        m = hashlib.sha3_384()
        m.update(secret)
        m.update(data)
        return m.digest()

    def test(self, secret):
        with open("cat2.jpeg","rb") as f:
            x = f.read()
        sHash = self.uploadData(x,secret)
        print(sHash.hex())
        y = self.getData(sHash)
        with open("res.jpeg","wb") as f:
            f.write(y)

class IotaFS():
    def __init__(self, token, fileCompression=False):
        self.api = Iota('https://nodes.thetangle.org:443', local_pow=True)
        self.blobStore = IotaFS_BlobStore(self.api)
        #self.token = token
        self.fileCompression = fileCompression
        self.hashState = hashlib.sha3_384()
        genesis = "This is the genesis block. lol."
        if self.fileCompression:
            raise Exception("Compression does not work currently")
            genesis += "#FILE COMPRESSION#"
        self.hashState.update(genesis.encode())
        self.hashState.update(token.encode())
        self._fileTree = {}
        self.lastBlockIncomplete = False
        self.incompleteBlockRescanTimeout = 5
        self.chainDelimiter = "#IOTA_FS#".encode()
        self.cache = {}

        self._fetchFileTree()

    def getFileTree(self, update=False):
        if update:
            self._fetchFileTree()
        return copy.deepcopy(self._fileTree)

    def _fetchFileTree(self):
        print("[<] Fetching FileTree")
        chain = bytes()
        while True:
            print("[<] Fetching FileTree-ChainBlock")
            sHash = self.hashState.digest()
            block = self._getBlob(sHash)
            data = block
            if self.fileCompression:
                data = gzip.decompress(data)
            if data==b'':
                print("[-] Last Block Received")
                break
            self.hashState.update(block)
            chain+=data
        if chain==b'':
            print("[.] FileTree succesfully fetched: [NO UPDATES]")
            return
        if chain.endswith(self.chainDelimiter):
            curRing = chain.split(self.chainDelimiter)[-2]
            self.lastBlockIncomplete = False
        else:
            print("[-] Last Block was incomplete; refetching...")
            self.lastBlockIncomplete = True
            time.sleep(self.incompleteBlockRescanTimeout)
            self._fetchFileTree()
            return
        print("{RING}: "+str(curRing))

        self._fileTree = msgpack.loads(curRing)
        print("[.] FileTree succesfully fetched: ")
        pprint(self._fileTree)

    def _mergeFileTrees(self, treeA, treeB):
        # We update treeB with values from treeA (treeA has priority), except for deletions,
        # which are always prioritized

        # fileTree = {fileA: sHash, fileB: sHash, dirA: {fileC: sHash}}
        for key, value in treeA.items():
            if isinstance(value, dict):
                # get node or create one
                node = treeB.setdefault(key, {})
                self._mergeFileTrees(value, node)
            else:
                if key in treeA and treeA[key]=="#REMOVE#" or key in treeB and treeB[key]=="#REMOVE#":
                    del treeB[key]
                treeB[key] = value

        return treeB

    def upsertFileTree(self, newFileTree):
        while self.lastBlockIncomplete:
            time.sleep(1)

        self._fileTree = self._mergeFileTrees(newFileTree, self._fileTree)

        newRing = msgpack.dumps(self._fileTree)+self.chainDelimiter
        sHash = self.hashState.digest()
        print("{RING}: "+str(newRing))
        payload = newRing
        if self.fileCompression:
            payload = gzip.compress(payload)
        self.blobStore.uploadDataRaw(payload, sHash)
        self.hashState.update(newRing) # For every link in the chain, we salt our hashState using the links data

    def _putBlob(self, data):
        return self.blobStore.uploadData(data, secrets.token_bytes(64))

    def _getBlob(self, sHash):
        data = self.blobStore.getData(sHash)
        return data

    def _fetchFile(self, sHash):
        file = self._getBlob(sHash)
        blob = file
        if self.fileCompression:
            blob = gzip.decompress(blob)
        #                    file  lastFetch    lastAccess
        self.cache[sHash] = [blob, time.time(), time.time()]
        return self.cache[sHash]

    def getFile(self, sHash):
        if sHash==b'':
            return [b'', 0, time.time()]
        print("/GET/ "+str(sHash)+" <- ")
        if sHash in self.cache and time.time()-self.cache[sHash][2] > MAX_CACHE_AGE:
            self.cache[sHash][2] = time.time()
            return self.cache[sHash]
        else:
            return self._fetchFile(sHash)

    def putFile(self, file, path):
        print("/PUT/ "+str(file)+" -> "+path)
        if file==b'':
            sHash = b''
        else:
            blob = file
            if self.fileCompression:
                blob = gzip.compress(blob)
            sHash = self._putBlob(blob)
            self.cache[sHash] = [file, time.time(), time.time()]
        treeDelta = {}
        subTree = treeDelta
        for elem in path.split("/")[:-1]:
            subTree[elem] = {}
            subTree = subTree[elem]
        subTree[path.split("/")[-1]] = sHash
        self.upsertFileTree(treeDelta)

    def mkdir(self, path):
        treeDelta = {}
        subTree = treeDelta
        for elem in path.split("/"):
            subTree[elem] = {}
            subTree = subTree[elem]
        self.upsertFileTree(treeDelta)
        return 0

    def removeFile(self, path):
        treeDelta = {}
        subTree = treeDelta
        for elem in path.split("/")[:-1]:
            subTree[elem] = {}
            subTree = subTree[elem]
        file = subTree[path.split("/")[-1]]
        subTree[path.split("/")[-1]] = "#REMOVE#"
        self.upsertFileTree(treeDelta)
        return file

    def mv(self, old, new):
        file = self.removeFile(old)
        self.putFile(self, file, new)

class IotaFS_Fuse(LoggingMixIn, Operations):
    def __init__(self, token, fileCompression=False):
        self.fs = IotaFS(token, fileCompression=fileCompression)

    def getSubtree(self, path):
        subTree = self.fs.getFileTree()
        for elem in path[1:].split("/"):
            if elem!="":
                if elem not in subTree:
                    return False
                if self.subtreeIsFile(subTree):
                    # we cannot traverse further, if this is a file
                    return False
                subTree = subTree[elem]
        return subTree

    def createFileObj(self, path, fileObj):
        subTree = self.fileTree
        for elem in path[1:].split("/")[:-1]:
            if elem not in subTree:
                return False
            subTree = subTree[elem]
        subTree[path.split("/")[-1]] = fileObj

    def subtreeIsFile(self, subtree):
        return isinstance(subtree, (bytes, bytearray))

    def subtreeExists(self, subtree):
        return not (subtree == False)

    def create(self, path, mode):
        print("[#] CREATE "+path)
        self.fs.putFile(b'', path[1:])
        #return open(path[1:])
        return 0

    #def destroy(self, path):
    #    self.sftp.close()
    #    self.client.close()

    def getattr(self, path, fh=None):
        print("[#] GETATTR "+path)
        subTree = self.getSubtree(path)
        if not self.subtreeExists(subTree):
            # File does not exist / is not a file
            raise FuseOSError(ENOENT)
        now = time.time()
        st = {}
        # mode decides access permissions and if file object is a directory (stat.S_IFDIR), file (stat.S_IFREG) or a special file
        if self.subtreeIsFile(subTree):
            st['st_mode']   = 0o744 | stat.S_IFREG
        else:
            st['st_mode']   = 0o744 | stat.S_IFDIR
        #st['st_ino']    = 0
        #st['st_dev']    = 0
        st['st_nlink']  = 1
        st['st_uid']    = os.getuid() #file object's user id
        st['st_gid']    = os.getgid() #file object's group id
        if fh and False:
            file, path, sHash, lastFetch, lastAccess = fh
            st["st_size"] = len(file)
            st['st_atime']  = lastAccess
            st['st_mtime']  = lastFetch
            st['st_ctime']  = 0
        else:
            st['st_size']   = len(sHash) + len(path) + 8 # Just an approximation...
            st['st_atime']  = 0 #last access time in seconds
            st['st_mtime']  = 0 #last modified time in seconds
            st['st_ctime']  = 0 # very old file
        st['st_blocks'] = 1 # Until we can skip blocks when reading / writing this should provide best possible performance...
        return st

    def mkdir(self, path, mode):
        print("[#] MKDIR "+path)
        self.fs.mkdir(path[1:])
        return 0

    def read(self, path, size, offset, fh):
        print("[#] READ "+path)
        file, path2, sHash, lastFetch, lastAccess = self.openFile(path)
        if path!=path2:
            print(path+"!="+path2)
            return FuseOSError(ENOENT)
        return file[offset : offset+size]

    def readdir(self, path, fh):
        print("[#] READDIR "+path)
        subTree = self.getSubtree(path)
        if self.subtreeIsFile(subTree):
            # We cant list a file!
            return FuseOSError(ENOENT)
        pprint(subTree)
        l = [".", ".."]
        for elem in subTree:
            l.append(elem)
        return l

    def rename(self, old, new):
        self.fs.mv(old,new)
        return 0

    def rmdir(self, path):
        self.fs.removeFile(path)
        return 0

    def write(self, path, data, offset, fh):
        print("[#] WRITE "+path)
        file, path2, sHash, lastFetch, lastAccess = self.openFile(path)
        if path!=path2:
            print(path+"!="+path2)
            return FuseOSError(ENOENT)

        raw = data
        file = file[:offset] + raw + file[offset+len(raw):]

        self.fs.putFile(file, path[1:])
        print("Write successfull")
        return len(raw)

    def openFile(self, path):
        subTree = self.getSubtree(path)
        if subTree == False:
            raise FuseOSError(ENOENT)
        else:
            if not self.subtreeIsFile(subTree):
                # cannot open a dir
                raise FuseOSError(ENOENT)
            sHash = subTree
            file, lastFetch, lastAccess = self.fs.getFile(sHash)
            return (file, path, sHash, lastFetch, lastAccess)

    def open(self, path, flags):
        return 0

    def release(self, path, fh):
        return 0

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('token')
    parser.add_argument('mount')
    args = parser.parse_args()

    fuse = FUSE(
        IotaFS_Fuse(args.token),
        args.mount,
        foreground=True,
        nothreads=True,
        allow_other=False)

#if __name__=="__main__":
#    iotaFS = IotaFS_BlobStore()
#
#    if len(sys.argv)>=2 and sys.argv[1]=="put":
#        print("Uploading '"+sys.argv[2]+"' using secret '"+" ".join(sys.argv[3:])+"'")
#        with open(sys.argv[2], "rb") as f:
#            x = f.read()
#        sHash = iotaFS.uploadData(x, " ".join(sys.argv[3:]).encode())
#        print("Stored at {"+sHash.hex()+"}")
#        print("Done.")
#    elif len(sys.argv)>=2 and sys.argv[1]=="get":
#        print("Downloading {"+sys.argv[2]+"} into '"+sys.argv[3]+"'")
#        with open(sys.argv[3], "wb") as f:
#            f.write(iotaFS.getData(bytearray.fromhex(sys.argv[2])))
#        print("Done.")
#    else:
#        print("Syntax:")
#        print("    put [file] [secret]")
#        print("    get [hash] [file]")
#