initial commit

main.py

@@ -0,0 +1,134 @@
+# 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]")