From eea1416911425a79293b970fa32657c8cd641e43 Mon Sep 17 00:00:00 2001 From: K4YT3X Date: Wed, 21 Aug 2019 11:03:37 -0400 Subject: [PATCH] organized code, using pathlib instead of string, using pyzipper instead of pyminizip, etc. --- prep_file.py | 161 ++++++++++++++++++++++++++++++--------------------- 1 file changed, 95 insertions(+), 66 deletions(-) diff --git a/prep_file.py b/prep_file.py index c7ca105..1292492 100644 --- a/prep_file.py +++ b/prep_file.py @@ -1,84 +1,113 @@ -#!/usr/bin/python +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Name: Prep File +Author: ytisf +Date of Creation: Unknown +Last Modified: May 26, 2019 -import os -import sys +Dev: K4YT3X +Last Modified: August 21, 2019 + +Licensed under the GNU General Public License Version 3 (GNU GPL v3), + available at: https://www.gnu.org/licenses/gpl-3.0.txt +(C) 2014-2019 ytisf +""" + +# built-in imports import hashlib +import pathlib +import sys +import time +import traceback try: - import pyminizip -except ImportError: - sys.stderr.write("Could not import 'pyminizip'. Did you install requirements?\n") - sys.stderr.write("You can always just get 'pyminizip' by 'pip install --user pyminizip'.\n") - sys.exit(1) + import pyzipper +except ImportError as e: + print('Could not import "pyzipper". Did you install requirements?', file=sys.stderr) + print('You can always just get "pyzipper" by "pip install --user pyzipper"', file=sys.stderr) + raise e -OUTPUT_FOLDER = "OUTPUT" +COMPRESSION_PASSWORD = 'infected' +OUTPUT_FOLDER = pathlib.Path('OUTPUT') -def _help(): +def print_help(): + """ print help message + + print program help message and return None """ - hmmmm. nope. - :return: - """ - print("Please run with '%s filename'." % sys.argv[0]) + print(f'usage: {__file__} [INPUT_FILE]') return -def _Do(file_path): +def prepare_file(file_path): + """ prep file from file path for submission + + take file name, encrypt in ZIP with password 'infected', create MD5 + and SHA1 sums and store all of that in a directory of it's own + + Arguments: + file_path {pathlib.Path} -- path object of input file """ - Prep file from file path for submission. Take file name, encrypt in ZIP with password 'infected', create MD5 - and SHA1 sums and store all of that in a directory of it's own. - :param file_path: str - :return: Bool - """ - if not os.path.isfile(file_path): - _help() - sys.stderr.write("Seems like '%s' is not a file.\n" % file_path) - return False + OUTPUT_FOLDER.mkdir(parents=True, exist_ok=True) - try: - os.mkdir(OUTPUT_FOLDER) - except OSError: - sys.stderr.write("Folder exists. Please remove it before continuing.\n") - return False + # create ZIP Archive + # we are using 7z because "zipfile" did not support adding a password + # Apparently "pyminizip" works just as well. + print('Info: Creating encrypted ZIP archive') + with pyzipper.AESZipFile(OUTPUT_FOLDER / f'{file_path.name}.zip', 'w', compression=pyzipper.ZIP_LZMA, encryption=pyzipper.WZ_AES) as zip_file: + zip_file.setpassword(COMPRESSION_PASSWORD.encode()) + zip_file.write(file_path) + print('Info: Created ZIP archive') - if "\\" in file_path: - filename = file_path.split("\\")[:-1] - elif "/" in file_path: - filename = file_path.split("/")[:-1] - else: - filename = file_path + # calculating file hashes + md5sum = hashlib.md5(open(OUTPUT_FOLDER / f'{file_path.name}.zip', 'rb').read()).hexdigest() + sha1sum = hashlib.sha1(open(OUTPUT_FOLDER / f'{file_path.name}.zip', 'rb').read()).hexdigest() - # Create ZIP Archive: - # We used 7z because 'zipfile' did not support adding a password. Apparently 'pyminizip' works just as well. - try: - pyminizip.compress(file_path, OUTPUT_FOLDER, "%s.zip" % filename, "infected", 9) - except Exception as e: - sys.stderr.write("Unknown error occurred. Please report this to us so that we can fix this.\n") - sys.stderr.write(str(e)) - return False - - compressed_path = '%s/%s.zip' % (OUTPUT_FOLDER, filename) - sys.stdout.write("[+]\tCreated ZIP Archive.\n") - md5sum = hashlib.md5(open(compressed_path, 'rb').read()).hexdigest() - sha1sum = hashlib.sha1(open(compressed_path, 'rb').read()).hexdigest() - open("%s/%s.md5" % (OUTPUT_FOLDER, filename), 'w').write(md5sum) - open("%s/%s.sha" % (OUTPUT_FOLDER, filename), 'w').write(sha1sum) - open("%s/%s.pass" % (OUTPUT_FOLDER, filename), 'w').write("infected") - return True + # writing file hashes and password to files + open(OUTPUT_FOLDER / f'{file_path.name}.md5', 'w').write(md5sum) + open(OUTPUT_FOLDER / f'{file_path.name}.sha', 'w').write(sha1sum) + open(OUTPUT_FOLDER / f'{file_path.name}.pass', 'w').write(COMPRESSION_PASSWORD) -if __name__ == "__main__": - if len(sys.argv) != 2: - _help() - sys.exit(1) - stt = _Do(sys.argv[1]) - if stt: - sys.stdout.write("Please don't forget to add details to 'conf/maldb.db' " - "and placing the folder in the appropriate directory.\n") - sys.stdout.write("Thanks for helping us get this accessible to everyone.\n") - sys.stdout.write("\n") - sys.exit(0) - else: - sys.exit(1) - +# start timer +start_time = time.time() + +# if this file is being imported +if __name__ != '__main__': + print('Error: This file cannot be imported', file=sys.stderr) + ImportError('File not importable') + +# check if there's a right amount of arguments provided +if len(sys.argv) != 2: + print_help() + exit(1) + +# convert input file path into file object +try: + input_file = pathlib.Path(sys.argv[1]) +except Exception: + print('Error: input file format invalid', file=sys.stderr) + +# input file validity check +if not input_file.is_file(): + print_help() + print(f'Seems like {str(input_file)} is not a file', file=sys.stderr) + exit(1) + +# zip file +try: + prepare_file(input_file) +except Exception: + print('Unexpected exception has been caught') + print('Compression has failed') + print('Please report the following error message to us so we can fix it') + traceback.print_exc() + exit(1) + +print('Script finished') +print(f'Time taken: {round((time.time() - start_time), 5)} seconds') +print('Please don\'t forget to add details to "conf/maldb.db" and placing the folder in the appropriate directory') +print('Thanks for helping us to get this accessible to everyone')