project.py

import hashlib
import math
import os
import re
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import tkinter as tk
from tkinter import filedialog

def get_aes_key_from_string(s: str, key_size=16) -> bytes:
    assert key_size in (16, 24, 32), "Invalid key size"
    hash_bytes = hashlib.sha256(s.encode()).digest()
    return hash_bytes[:key_size]

def encrypt(plainbytes: bytes, key: bytes) -> bytes:
    cipher = AES.new(key, AES.MODE_ECB)
    padded = pad(plainbytes, AES.block_size)
    return cipher.encrypt(padded)

def decrypt(ciphertext: bytes, key: bytes) -> bytes:
    cipher = AES.new(key, AES.MODE_ECB)
    decrypted = unpad(cipher.decrypt(ciphertext), AES.block_size)
    return decrypted

def encode_file(full_path: str, entered_key: str):
    with open(full_path, "rb") as input_file:
        file_bytes = input_file.read()

    encoded_bytes = encrypt(file_bytes, get_aes_key_from_string(entered_key))

    encoded_file_path = full_path+".encoded"
    with open(encoded_file_path, "wb") as encoded_file:
        encoded_file.write("<<<EncodedWithTrampoline>>>".encode() + encoded_bytes + "<<<EncodedWithTrampoline>>>".encode())
    
    return encoded_file_path

def decode_file(full_path: str, entered_key: str):
    with open(full_path, "rb") as encoded_file:
        encoded_bytes = encoded_file.read()

    encoded_bytes = encoded_bytes.removeprefix(b"<<<EncodedWithTrampoline>>>")
    encoded_bytes = encoded_bytes.removesuffix(b"<<<EncodedWithTrampoline>>>")

    decoded_bytes = decrypt(encoded_bytes, get_aes_key_from_string(entered_key))

    directory = os.path.dirname(full_path)
    file_name = os.path.basename(full_path)

    file_name = "decoded_" + file_name.removesuffix(".encoded")

    new_path = os.path.join(directory, file_name)

    with open(new_path, "wb") as decoded_file:
        decoded_file.write(decoded_bytes)
    
    return new_path

def get_list_of_encoded_files(directory_path):

    list_of_files = []
    for entry in os.listdir(directory_path):
        full_path = os.path.join(directory_path, entry)
        if os.path.isfile(full_path):
            if check_validity(full_path):
                list_of_files.append(full_path)

    return list_of_files

def check_validity(full_path):

    with open(full_path, "rb") as encoded_file:
        encoded_bytes = encoded_file.read()

    if encoded_bytes.startswith(b"<<<EncodedWithTrampoline>>>") and encoded_bytes.endswith(b"<<<EncodedWithTrampoline>>>"):
        return True
    
    return False

def file_splitter(full_path, parts_count):

    with open(full_path, "rb") as f:
        file_content = f.read()
    
    original_file_dir = os.path.dirname(full_path)
    original_file_name = os.path.basename(full_path)

    file_extension = original_file_name.split(".")[-1]

    file_hash = hashlib.sha256(file_content).hexdigest()

    part_size = math.ceil(len(file_content) / parts_count)
    
    file_paths = []
    for i in range(parts_count):
        fname = original_file_name + f".part{i+1}"

        with open(fname, "wb") as f:
            f.write(f"[part{i+1}-{file_extension}-{file_hash}]".encode())
            f.write(file_content[i*part_size:(i+1)*part_size])
        
        file_paths.append(os.path.join(original_file_dir, fname))
    
    return file_paths

def file_joiner(file_paths):

    try:
        invalid = False
        fcontents = []
        for fpath in file_paths:
            with open(fpath, "rb") as f:
                fcontent = f.read()
            if not fcontent.startswith(b"[part"):
                invalid = True
                break
            fcontents.append(fcontent)
        
        if not invalid:
            fcontents = sorted(fcontents, key=lambda x:x[:6])

            raw_contents = []
            desc_hash = None
            file_extension = None

            for fcontent in fcontents:
                pattern = re.compile(rb"^\[part.*?\]")

                match = pattern.match(fcontent)
                if match:
                    desc = match.group(0)
                    desc = desc.decode()

                    if not file_extension:
                        file_extension = desc.split("-")[1]
                    elif file_extension != desc.split("-")[1]:
                        invalid = True
                        break

                    if not desc_hash:
                        desc_hash = desc.split("-")[-1][:-1]
                    elif desc_hash != desc.split("-")[-1][:-1]:
                        invalid = True
                        break

                    raw_content = fcontent[match.end():]
                    raw_contents.append(raw_content)
                else:
                    invalid = True
                    break
    
        if invalid:
            print("One or more files from the files you provided is not a split file generated by me or the files aren't generated from the same file!")
        else:
            final_content = b"".join(raw_contents)

            if hashlib.sha256(final_content).hexdigest() == desc_hash:
                
                file_dir = os.path.dirname(file_paths[0])
                new_path = os.path.join(file_dir, f"file-{desc_hash}.{file_extension}")
                with open(new_path, "wb") as f:
                    f.write(final_content)
                

                return new_path, final_content
            else:
                print("Unfortunately the stored hash does not match the result file's hash")
    except FileNotFoundError as ex:
        print("One or more of the paths you passed wasn't a valid file path")
    except Exception as ex:
        print("One of the files wasn't how is supposed to be")


def main():

    root = tk.Tk()
    root.withdraw()

    print("Welcome to Trampoline file encoder :)")

    while True:
        inp = input("What do you want to do?" \
            "\n\t1: Encode a file" \
            "\n\t2: List of encoded files in a directory" \
            "\n\t3: Decode a file" \
            "\n\t4: Split a file" \
            "\n\t5: Rejoin files" \
            "\n\t0: Exit" \
            "\n: "
        )

        if inp == "1":
            full_path = input("Enter the file full path: ")
            entered_key = input("Enter a key to encrypt the file with: ")

            try:
                encryption_path = encode_file(full_path, entered_key)
                print("Encryption Done")
                print("Encrypted file:", encryption_path)
            except FileNotFoundError as ex:
                print("\nError: The path you entered wasn't a valid path to a file!")
                
        
        elif inp == "2":
            directory_path = input("Enter the full path of the directory: ")

            try:
                list_of_files = get_list_of_encoded_files(directory_path)

                print()

                for f in list_of_files:
                    print(f, "is encoded")
                
                if not list_of_files:
                    print("No encoded files are found in this directory")
            except FileNotFoundError as ex:
                print("\nErro: Not a valid path")
            
        elif inp == "3":
            full_path = input("Enter the full path of the encoded file: ")
            entered_key = input("Enter the key you used to encode: ")

            try:
                if check_validity(full_path):

                    decryption_path = decode_file(full_path, entered_key)
                    
                    print("Decryption Done")
                    print("Decrypted file:", decryption_path)
                else:
                    print("Not a valid encoded file!")
            except FileNotFoundError as ex:
                print("\nError: Not a valid path to a file")
        
        elif inp == "4":
            try:
                full_path = input("Enter the full path of the file: ")
                parts_count = input("Enter the number of parts you want it to be split to (max 8): ")

                parts_count = min(8, int(parts_count))

                file_paths = file_splitter(full_path, parts_count)
                
                print("Parts' path:")
                print("\n".join(file_paths))
            except ValueError as ex:
                print("The part count you entered wasn't a number")
            except FileNotFoundError as ex:
                print("\nError: Not a valid path to a file")
            
        elif inp == "5":
            root.update()
            root.attributes('-topmost', True)   # force dialogs on top
            root.lift()

            file_paths = filedialog.askopenfilenames(
                title="Select files",
                filetypes=[("All Files", "*.*"), ("Text Files", "*.txt")]
            )

            if file_paths:
                new_path, _ = file_joiner(file_paths)
                print("Your original file was rejoined to this path: ", new_path)
            
        elif inp == "0":
            print("Exiting the program.")
            exit()
        
        else:
            print("Invalid input")
        
        print("\n")


if __name__ == "__main__":
    main()