shamir/main.py

import argparse
from os.path import join, exists, isfile, basename
from os import makedirs, listdir

from core import shamir_decode, shamir_encode


def main():
    parser = argparse.ArgumentParser("Shamir's Secret Sharing")
    parser.set_defaults(mode='None')

    subparsers = parser.add_subparsers()

    split = subparsers.add_parser("split")
    split.add_argument("file", help="File to split into multiple files")
    split.add_argument("-n", type=int, required=True,
                       help="Number of parts to split the file.")
    split.add_argument("-k", type=int, required=True,
                       help="Expected number of parts needed to recover the original file.")
    split.add_argument("-o", "--output", default='output',
                       help="Output directory of all the files.")
    split.set_defaults(mode='split')

    recover = subparsers.add_parser("recover")
    recover.add_argument("directory")
    recover.add_argument("-p", "--prefix", default="",
                         help="Read parts from files in directory with this prefix.")
    recover.add_argument("-o", "--output", default='',
                         help="Output file to save the recovered file.")
    recover.set_defaults(mode='recover')

    args = parser.parse_args()

    if args.mode == 'split':
        name = basename(args.file)

        for part_id in range(args.n):
            name_file = f'{name}.{part_id+1}.shamir'
            path = join(args.output, name_file)
            if exists(path):
                print(f"File: {path} already exists, use another directory.")
                exit(1)

        makedirs(args.output, exist_ok=True)

        with open(args.file, 'rb') as f:
            data = f.read()

        parts = shamir_encode(data, args.k, args.n)

        for part_id, part in enumerate(parts):
            name_file = f'{name}.{part_id+1}.shamir'
            path = join(args.output, name_file)
            with open(path, 'wb') as f:
                f.write(part)

    elif args.mode == 'recover':
        parts = []
        expected_name = False

        for name in listdir(args.directory):
            if name.endswith('.shamir') and name.startswith(args.prefix):
                file_name = join(args.directory, name)

                if isfile(file_name) and name.endswith('shamir'):
                    expected_name = '.'.join(name.split('.')[:-2])
                    with open(file_name, 'rb') as f:
                        parts.append(f.read())

        if len(parts) == 0:
            print("No parts found")
            exit(2)

        makedirs(args.output, exist_ok=True)
        output = join(args.output, expected_name)

        if exists(output):
            print(f"File: {output} already exists, use another file.")

        secret = shamir_decode(parts)

        with open(output, 'wb') as f:
            f.write(secret)
    else:
        parser.print_help()


if __name__ == '__main__':
    main()