Browse Source

Implement CLI

Marcelo Fornet 4 năm trước cách đây
mục cha
commit
a14002b61d
2 tập tin đã thay đổi với 43 bổ sung5 xóa
  1. 2 0
      .gitignore
  2. 41 5
      main.py

+ 2 - 0
.gitignore

@@ -1,2 +1,4 @@
 __pycache__/
 .vscode/
+output/
+build/

+ 41 - 5
main.py

@@ -1,5 +1,6 @@
 import argparse
-from os.path import join
+from os.path import join, exists, isfile
+from os import makedirs, listdir
 
 from core import shamir_decode, shamir_encode
 
@@ -16,23 +17,33 @@ def main():
                        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=".",
+    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",
+    recover.add_argument("-p", "--prefix", default="",
                          help="Read parts from files in directory with this prefix.")
-    recover.add_argument("-o", "--output",
+    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':
+        for part_id in range(args.n):
+            name_file = f'{args.file}.{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):
@@ -42,7 +53,32 @@ def main():
                 f.write(part)
 
     elif args.mode == 'recover':
-        pass
+        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()