# 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 time
import sys

api = Iota('https://nodes.thetangle.org:443', local_pow=True)

def genBundles(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(
            api.prepare_transfer(
                transfers = bundleTxs,
                inputs = [addr]
            )['trytes']
        )
    return bundles

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

def uploadData(data, secret):
    print("Uploading...")
    m = hashlib.sha3_384()
    m.update(secret)
    m.update(data)
    sHash = m.digest()
    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 = genBundles(ct_bytes, addrIter)
    sendBundles(bundles)
    return sHash

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

def getData(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 = api.find_transactions(tags=[Tag("IOTAFS")], addresses=[addr])["hashes"]
        if len(txHash)==0:
            break
        bundles = api.get_bundles(txHash[0])["bundles"]
        for bundle in bundles:
            for tx in bundle.transactions:
                tryteMsg+=str(tx.signature_message_fragment)
    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(sHash):
    return getData(sHash).decode("utf-8")

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

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

if __name__=="__main__":
    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 = 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(getData(bytearray.fromhex(sys.argv[2])))
        print("Done.")
    else:
        print("Syntax:")
        print("    put [file] [secret]")
        print("    get [hash] [file]")