|
@@ -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()
|
|
|
|