#!/usr/bin/python import zfec from Crypto.Cipher import AES from Crypto import Random import sys def tmapzfecsplit(data, m, k): splitter = zfec.easyfec.Encoder(k, m) splitdata = splitter.encode(data) # calc pad padlen = len(splitdata[0])*k - len(data) # prepend k,m sharenum, and padlen to each block - can't recreate without it splitdata = [str(k)+","+str(m)+","+str(snum)+","+str(padlen)+":"+sdata for snum, sdata in enumerate(splitdata)] return splitdata def tmaprecreate(datablocks): sharenums = [] cleanblocks = [] k = 0 m = 0 padlen = 0 for block in datablocks: splitdata = block.split(":", 1) cleanblocks.append(splitdata[1]) metadata = splitdata[0].split(",") k = int(metadata[0]) m = int(metadata[1]) sharenums.append(int(metadata[2])) padlen = int(metadata[3]) if len(cleanblocks) < k: # not enough blocks provided raise Exception("Can't recreate file. Need {} parts, only have {}.".format(k, len(cleanblocks))) decoder = zfec.easyfec.Decoder(k,m) origdata = decoder.decode(cleanblocks, sharenums, padlen) return origdata def tmap_encryptdata(data): # AES-256 encrypts the data, pre-pending the IV and appending the key # the IV is the stringified padding value to get to 16 byte blocksize # 32 byte key is random padlen = 16 - (len(data) % 16) iv = "0" * (16- len(str(padlen))) + str(padlen) data = data + "0" * padlen key = Random.new().read(32) cipher = AES.new(key, AES.MODE_CBC, iv) encryptedData = iv + cipher.encrypt(data) + key return encryptedData def tmap_decryptdata(encryptedData): # AES-256 decrypts data # expects first 16 bytes is IV (and padlen) # and last 32 is the key iv = encryptedData[:16] key = encryptedData[-32:] cipher = AES.new(key, AES.MODE_CBC, iv) msg = encryptedData[16:-32] data = cipher.decrypt(msg) padlen = int(iv) data = data[:-padlen] return data def tmap_split(filename, m, k): # encrypts data in filename # splits the data into m parts, k of which can recreate the file # prepends the filename to each file # saves as filename-X.tmap f = open(filename) data = f.read() f.close() enc_data = tmap_encryptdata(data) datablocks = tmapzfecsplit(enc_data, m, k) for filenum, data in enumerate(datablocks): f = open("{}-{}.tmap".format(filename, filenum+1), "wb") data = filename+":"+data f.write(data) f.close() return def tmap_rebuild(filenames): # rebuilds a tmap file from the parts in filenames savefilename = "" datablocks = [] for name in filenames: fr = open(name, "rb") data_raw = fr.read() fr.close() data_split = data_raw.split(":", 1) if savefilename == "": savefilename = data_split[0] elif not savefilename == data_split[0]: raise Exception("Not all data is from the same file. Cannot recreate") datablocks.append(data_split[1]) enc_data = tmaprecreate(datablocks) data = tmap_decryptdata(enc_data) f = open(savefilename, "wb") f.write(data) def tmap_cmdline(args): # not using getopt to save code space, but harder to detect errors help_text = 'Usage:\ntmap --make m k filename\n\tSplit filename into m parts, k of which are needed to recover\n' help_text = help_text + 'tmap --recover file1 file2 file3 ... filen\n\tAttempt to recover the file from the listed parts.' help_text = help_text + '\nArguments must be in order!!' if len(args) < 3: print(help_text) sys.exit(2) try: if args[1] == "--make": m = int(args[2]) k = int(args[3]) tmap_split(args[4], m, k) print("{} split into {} parts, {} needed to recover.".format(args[4], m, k)) elif args[1] == "--recover": tmap_rebuild(args[2:]) print("File recovered!") else: print(help_text) except Exception as e: print("Failed to work:") print(e.args[0]) sys.exit(2) tmap_cmdline(sys.argv)