From 9d5680e360209a45374f0640ba50aaa2f35d5e53 Mon Sep 17 00:00:00 2001 From: AlfredoSequeida Date: Fri, 19 Feb 2021 17:07:36 -0800 Subject: [PATCH] documented code, formatted using PEP 8 style guide, and made use of type hinting --- fvid/fvid.py | 305 ++++++++++++++++++++++++++++++++++++--------------- 1 file changed, 215 insertions(+), 90 deletions(-) diff --git a/fvid/fvid.py b/fvid/fvid.py index ddd1d18..3857078 100644 --- a/fvid/fvid.py +++ b/fvid/fvid.py @@ -18,70 +18,100 @@ from Crypto.Cipher import AES try: from fvid_cython import cy_get_bits_from_image as cy_gbfi + use_cython = True except (ImportError, ModuleNotFoundError): use_cython = False FRAMES_DIR = "./fvid_frames/" -SALT = '63929291bca3c602de64352a4d4bfe69'.encode() # It need be the same in one instance of coding/decoding -DEFAULT_KEY = ' '*32 +SALT = ( + "63929291bca3c602de64352a4d4bfe69".encode() +) # It needs be the same in one instance of coding/decoding +DEFAULT_KEY = " " * 32 DEFAULT_KEY = DEFAULT_KEY.encode() NOTDEBUG = True -TEMPVIDEO = '_temp.mp4' -FRAMERATE = '1' +TEMPVIDEO = "_temp.mp4" +FRAMERATE = "1" + class WrongPassword(Exception): pass + class MissingArgument(Exception): pass -def get_password(password_provided): - if password_provided=='default': + +def get_password(password_provided: str) -> bytes: + """ + Prompt user for password and create a key for decrypting/encrypting + + password_provided: password provided by tge user with -p flag + """ + + if password_provided == "default": return DEFAULT_KEY else: if password_provided == None: password_provided = getpass.getpass("Enter password:") - password = str(password_provided).encode() + password = str(password_provided).encode() kdf = PBKDF2HMAC( algorithm=hashes.SHA512(), length=32, salt=SALT, iterations=100000, - backend=default_backend() - ) + backend=default_backend(), + ) key = kdf.derive(password) return key -def get_bits_from_file(filepath, key): - print('Reading file...') + +def get_bits_from_file(filepath: str, key: bytes) -> BitArray: + """ + Get/read bits from file + + filepath -- the file to read + key -- key used to encrypt file + """ + + print("Reading file...") + bitarray = BitArray(filename=filepath) + # adding a delimiter to know when the file ends to avoid corrupted files # when retrieving - cipher = AES.new(key, AES.MODE_EAX, nonce=SALT) ciphertext, tag = cipher.encrypt_and_digest(bitarray.tobytes()) - + filename = os.path.basename(filepath) - pickled = pickle.dumps({'tag':tag, - 'data':ciphertext, - 'filename':filepath}) - print('Zipping...') - #zip + pickled = pickle.dumps( + {"tag": tag, "data": ciphertext, "filename": filepath} + ) + + print("Zipping...") + + # zip out = io.BytesIO() - with gzip.GzipFile(fileobj=out, mode='w') as fo: + with gzip.GzipFile(fileobj=out, mode="w") as fo: fo.write(pickled) zip = out.getvalue() - #zip - + # zip + del bitarray del pickled bitarray = BitArray(zip) + return bitarray.bin -def get_bits_from_image(image): + +def get_bits_from_image(image: Image) -> str: + """ + extract bits from image (frame) pixels + + image -- png image file used to extract bits from + """ if use_cython: bits = cy_gbfi(image) @@ -91,7 +121,7 @@ def get_bits_from_image(image): px = image.load() bits = "" - + for y in range(height): for x in range(width): @@ -99,9 +129,13 @@ def get_bits_from_image(image): pixel_bin_rep = "0" - # if the white difference is smaller, that means the pixel is closer - # to white, otherwise, the pixel must be black - if abs(pixel[0] - 255) < abs(pixel[0] - 0) and abs(pixel[1] - 255) < abs(pixel[1] - 0) and abs(pixel[2] - 255) < abs(pixel[2] - 0): + # if the white difference is smaller, that means the pixel is + # closer to white, otherwise, the pixel must be black + if ( + abs(pixel[0] - 255) < abs(pixel[0] - 0) + and abs(pixel[1] - 255) < abs(pixel[1] - 0) + and abs(pixel[2] - 255) < abs(pixel[2] - 0) + ): pixel_bin_rep = "1" # adding bits @@ -110,93 +144,139 @@ def get_bits_from_image(image): return bits -def get_bits_from_video(video_filepath): - # get image sequence from video - print('Reading video...') +def get_bits_from_video(video_filepath: str) -> str: + """ + extract the bits from a video by frame (using a sequence of images) + + video_filepath -- The file path for the video to extract bits from + """ + + print("Reading video...") + image_sequence = [] - os.system('ffmpeg -i ' + video_filepath + ' -c:v libx264rgb -filter:v fps=fps=' + FRAMERATE + ' ' + TEMPVIDEO) - os.system('ffmpeg -i ' + TEMPVIDEO + ' ./fvid_frames/decoded_frames_%d.png') - os.remove(TEMPVIDEO) - for filename in sorted(glob.glob(f"{FRAMES_DIR}decoded_frames*.png"), key=os.path.getmtime): + os.system( + "ffmpeg -i " + + video_filepath + + " -c:v libx264rgb -filter:v fps=fps=" + + FRAMERATE + + " " + + TEMPVIDEO + ) + os.system( + "ffmpeg -i " + TEMPVIDEO + " ./fvid_frames/decoded_frames_%d.png" + ) + os.remove(TEMPVIDEO) + + for filename in sorted( + glob.glob(f"{FRAMES_DIR}decoded_frames*.png"), key=os.path.getmtime + ): image_sequence.append(Image.open(filename)) bits = "" sequence_length = len(image_sequence) - print('Bits are in place') + print("Bits are in place") + if use_cython: - print('Using Cython...') + print("Using Cython...") for index in tqdm(range(sequence_length)): bits += get_bits_from_image(image_sequence[index]) return bits -def save_bits_to_file(file_path, bits, key): - # get file extension +def save_bits_to_file(file_path: str, bits: str, key: bytes): + """ + save/write bits to a file + + file_path -- the path to write to + bits -- the bits to write + key -- key userd for file decryption + """ + + # get file extension bitstring = Bits(bin=bits) - #zip - print('Unziping...') + # zip + print("Unziping...") in_ = io.BytesIO() in_.write(bitstring.bytes) in_.seek(0) - with gzip.GzipFile(fileobj=in_, mode='rb') as fo: + with gzip.GzipFile(fileobj=in_, mode="rb") as fo: bitstring = fo.read() - #zip - + # zip unpickled = pickle.loads(bitstring) - tag = unpickled['tag'] - ciphertext = unpickled['data'] - filename = unpickled['filename'] - + tag = unpickled["tag"] + ciphertext = unpickled["data"] + filename = unpickled["filename"] + cipher = AES.new(key, AES.MODE_EAX, nonce=SALT) bitstring = cipher.decrypt(ciphertext) - print('Checking integrity...') + + print("Checking integrity...") + try: - cipher.verify(tag) + cipher.verify(tag) except ValueError: - raise WrongPassword("Key incorrect or message corrupted") + raise WrongPassword("Key incorrect or message corrupted") bitstring = BitArray(bitstring) - - # If filepath not passed in use defualt - # otherwise used passed in filepath + # If filepath not passed in use default otherwise used passed in filepath if file_path == None: filepath = filename else: - filepath = file_path # No need for mime Magic + filepath = file_path - with open( - filepath, "wb" - ) as f: + with open(filepath, "wb") as f: bitstring.tofile(f) +def split_string_by_n(bitstring: str, n: int) -> list: + """ + Split a string every n number of characters + (or less if the 'remaining characters' < n ) this way we can sperate the + data for an etire video into a list based on the resolution of a frame. -def split_list_by_n(lst, n): - for i in range(0, len(lst), n): - yield lst[i : i + n] + bitstring -- a string containing bits + n -- split the string every n characters, for example to split a + 1920 x 1080 frame, this would be 1920*1080 = 2073600 + """ + + bit_list = [] + + for i in range(0, len(bitstring), n): + bit_list.append(bitstring[i : i + n]) + + return bit_list + + +def make_image_sequence(bitstring: BitArray, resolution: tuple = (1920, 1080)): + """ + Create image sequence (frames) for a video + + bitstring -- BitArray of bits used to create pixels with bit data + resolution -- the resoultion used for each frame (default 1920x1080) + """ -def make_image_sequence(bitstring, resolution=(1920, 1080)): width, height = resolution # split bits into sets of width*height to make (1) image set_size = width * height # bit_sequence = [] - print('Making image sequence') - print('Cutting...') - bitlist = list(split_list_by_n(bitstring, set_size)) - + print("Making image sequence") + print("Cutting...") + + bitlist = split_string_by_n(bitstring, set_size) + del bitstring - - bitlist[-1] = bitlist[-1] + '0'*(set_size - len(bitlist[-1])) + + bitlist[-1] = bitlist[-1] + "0" * (set_size - len(bitlist[-1])) index = 1 bitlist = bitlist[::-1] - print('Saving frames...') + print("Saving frames...") for _ in tqdm(range(len(bitlist))): bitl = bitlist.pop() image_bits = list(map(int, bitl)) @@ -204,35 +284,49 @@ def make_image_sequence(bitstring, resolution=(1920, 1080)): image = Image.new("1", (width, height)) image.putdata(image_bits) - image.save( - f"{FRAMES_DIR}encoded_frames_{index}.png" - ) + image.save(f"{FRAMES_DIR}encoded_frames_{index}.png") index += 1 -def make_video(output_filepath, framerate=FRAMERATE): + +def make_video(output_filepath: str, framerate: int = FRAMERATE): + """ + Create video using ffmpeg + + output_filepath -- the output file path where to store the video + framerate -- the framerate for the vidoe (default 1) + """ if output_filepath == None: outputfile = "file.mp4" else: outputfile = output_filepath - os.system('ffmpeg -r ' + framerate + ' -i ./fvid_frames/encoded_frames_%d.png -c:v libx264rgb ' + outputfile) - + os.system( + "ffmpeg -r " + + framerate + + " -i ./fvid_frames/encoded_frames_%d.png -c:v libx264rgb " + + outputfile + ) def cleanup(): - # remove frames + """ + Clean up the files (frames) creted by fvid during encoding/decoding + """ import shutil shutil.rmtree(FRAMES_DIR) def setup(): - import os + """ + setup fvid directory used to store frames for encoding/decoding + """ if not os.path.exists(FRAMES_DIR): os.makedirs(FRAMES_DIR) + def main(): global FRAMERATE parser = argparse.ArgumentParser(description="save files as videos") @@ -245,24 +339,47 @@ def main(): parser.add_argument("-i", "--input", help="input file", required=True) parser.add_argument("-o", "--output", help="output path") - parser.add_argument("-f", "--framerate", help="set framerate for encoding (as a fraction)", default=FRAMERATE, type=str) - parser.add_argument("-p", "--password", help="set password", nargs="?", type=str, default='default') + parser.add_argument( + "-f", + "--framerate", + help="set framerate for encoding (as a fraction)", + default=FRAMERATE, + type=str, + ) + parser.add_argument( + "-p", + "--password", + help="set password", + nargs="?", + type=str, + default="default", + ) args = parser.parse_args() setup() - if not NOTDEBUG: - print('args', args) - print('PASSWORD', args.password, [len(args.password) if len(args.password) is not None else None for _ in range(0)]) + if not NOTDEBUG: + print("args", args) + print( + "PASSWORD", + args.password, + [ + len(args.password) if len(args.password) is not None else None + for _ in range(0) + ], + ) + + # using default framerate if none is provided by the user if args.framerate != FRAMERATE: FRAMERATE = args.framerate + # check for arguments if not args.decode and not args.encode: - raise MissingArgument('You should use either --encode or --decode!') #check for arguments - + raise MissingArgument("You should use either --encode or --decode!") + key = get_password(args.password) - + if args.decode: bits = get_bits_from_video(args.input) @@ -274,22 +391,30 @@ def main(): save_bits_to_file(file_path, bits, key) elif args.encode: - # isdigit has the benefit of being True and raising an error if the user passes a negative string - # all() lets us check if both the negative sign and forward slash are in the string, to prevent negative fractions - if (not args.framerate.isdigit() and "/" not in args.framerate) or all(x in args.framerate for x in ("-", "/")): - raise NotImplementedError("The framerate must be a positive fraction or an integer for now, like 3, '1/3', or '1/5'!") + + # isdigit has the benefit of being True and raising an error if the + # user passes a negative string + # all() lets us check if both the negative sign and forward slash are + # in the string, to prevent negative fractions + if (not args.framerate.isdigit() and "/" not in args.framerate) or all( + x in args.framerate for x in ("-", "/") + ): + raise NotImplementedError( + "The framerate must be a positive fraction or an integer for " + "now, like 3, '1/3', or '1/5'!" + ) + # get bits from file bits = get_bits_from_file(args.input, key) # create image sequence make_image_sequence(bits) - video_file_path = None if args.output: video_file_path = args.output make_video(video_file_path, args.framerate) - - cleanup() + + cleanup() \ No newline at end of file