# iotaFS a.k.a iotaShitPoc

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

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 ""
        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())
        #sHash = getSHash(x, "catSecret".encode())
        y = self.getData(sHash)
        with open("res.jpeg","wb") as f:
            f.write(y)

class IotaFS():
    def __init__(self, token):
        self.api = Iota('https://nodes.thetangle.org:443', local_pow=True)
        self.blobStore = IotaFS_BlobStore(self.api)
        #self.token = token
        self.hashState = hashlib.sha3_384()
        self.hashState.update(token.encode())
        self._fileTree = {}
        self.lastBlockIncomplete = False
        self.incompleteBlockRescanTimeout = 5
        self.chainDelimiter = "#CHAIN_DELIM#"
        self.cache = {}

        self._fetchFileTree()

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

    async def fileTreeFetchLoop(self, interval=10):
        while True:
            self._fetchFileTree()
            await asyncio.sleep(10)

    def _fetchFileTree(self):
        print("[<] Fetching FileTree")
        chain = ""
        while True:
            print("[<] Fetching FileTree-ChainBlock")
            sHash = self.hashState.digest()
            block = self._getBlob(sHash)
            if block=="":
                print("[-] Last Block Received")
                break
            self.hashState.update(block)
            chain+=block
        if chain=="":
            print("[.] FileTree succesfully fetched: [EMPTY FILE TREE]")
            return
        if chain.endswith(self.chainDelimiter):
            curRing = chain.split(self.chainDelimiter)[-1]
            self.lastBlockIncomplete = False
        else:
            print("[-] Last Block was incomplete; refetching...")
            self.lastBlockIncomplete = True
            time.sleep(self.incompleteBlockRescanTimeout)
            self._fetchFileTree()

        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#":
                    treeB[key] = "#REMOVE#"
                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()
        self.blobStore.uploadDataRaw(newRing.encode(), sHash)
        self.hashState.update(newRing) # For every link in the chain, we salt our hashState using the links data

    def _putBlob(self, data):
        #TODO: Use secure random provider
        return self.blobStore.uploadData(data, str(random.random()*999999999999).encode())

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

    def _fetchFile(self, sHash):
        file = self._getBlob(sHash)
        #                    file  lastFetch   lastAccess
        self.cache[sHash] = [file, time.now(), time.now()]
        return self.cache[sHash]

    def getFile(self, sHash):
        if sHash in self.cache:
            # TODO: maybe update if to old?
            self.cache[sHash][2] = time.now()
            return self.cache[sHash]
        else:
            return self._fetchFile(sHash)

    def putFile(self, file, path):
        blablabla


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

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

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

    def subtreeIsFile(self, subtree):
        return isinstance(subtree, str)

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

    #def chmod(self, path, mode):
    #    return self.sftp.chmod(path, mode)

    #def chown(self, path, uid, gid):
    #    return self.sftp.chown(path, uid, gid)

    #def create(self, path, mode):
    #    f = self.sftp.open(path, 'w')
    #    f.chmod(mode)
    #    f.close()
    #    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']   = 777 | stat.S_IFREG
        else:
            st['st_mode']   = 777 | 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:
            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']   = 0 # 0 Byte lol
            st['st_atime']  = now  #last access time in seconds
            st['st_mtime']  = now  #last modified time in seconds
            st['st_ctime']  = 0 # very old file
        # TODO: Actuall real value
        block_size = 512
        st['st_blocks'] = (int) ((st['st_size'] + block_size-1) / block_size)
        return st

    #def mkdir(self, path, mode):
    #    return self.sftp.mkdir(path, mode)

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

        return file[offset : offset+size]

    def readdir(self, path, fh):
        print("[#] READDIR "+path)
        l = ['.', '..']
        for elem in self.getSubtree(path):
            l.append(elem.encode('utf-8'))
        return l

    #def rename(self, old, new):
    #    return self.sftp.rename(old, new)

    #def rmdir(self, path):
    #    return self.sftp.rmdir(path)

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

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

    def open(self, path, flags):
        print("[#] OPEN "+path)
        subTree = self.getSubtree(path)
        if subTree == False:
            pass
        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)

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]")
#