diff --git a/rename_viruses.py b/rename_viruses.py index 3b25f8f..333d731 100644 --- a/rename_viruses.py +++ b/rename_viruses.py @@ -4,93 +4,104 @@ import os import sys import re -my_parser = argparse.ArgumentParser(description='Rename all suspicious files found by clamav') - -my_parser.add_argument('File', - metavar='file', - type=str, - help='path to clamav log file') - -my_parser.add_argument('-s', +# Setup argparse +argparser = argparse.ArgumentParser(description='Rename files identified as dangerous by clamav') +argparser.add_argument('-f', + '--file', + help='Path to clamav log file', + type=argparse.FileType('r', encoding='UTF-8'), + default=(None if sys.stdin.isatty() else sys.stdin)) +argparser.add_argument('-s', '--suffix', type=str, - help='suffix to add to the end of the file names', + help='Suffix to be appended to the end of the filenames', default='.VIRUS') - -my_parser.add_argument('-u', +argparser.add_argument('-u', '--undo', action='store_true', - help='undo renaming') - -my_parser.add_argument('--dry-run', + help='Undo renaming') +argparser.add_argument('-d', + '--dry-run', action='store_true', - help='perform a test run where no file names are actually changed') + help='Perform a test run where no file names are actually changed') def main(): - args = my_parser.parse_args() - file = args.File + args = argparser.parse_args() success_count = 0 error_count = 0 warning_count = 0 - if not os.path.isfile(file): - print(f"The file '{file}' does not exist") + # Print info if no file name is specified or passed via stdin + if not args.file: + print('Please specify a path for a clamav log file using the -f argument or by piping the output directly into ' + 'this program.') sys.exit() - with open(file) as f: - lines = f.readlines() - files_to_rename = [re.split(r":", line.strip())[0] for line in lines if re.search(r"FOUND$", line)] + # Read input file + with args.file as file: + lines = file.readlines() - for file_to_rename in files_to_rename: - new_name = file_to_rename + args.suffix - if not args.undo: - if not os.path.isfile(file_to_rename): - if os.path.isfile(new_name): - print(f"WARNING: The file {file_to_rename} was already renamed to {new_name}") - warning_count += 1 - else: - print(f"ERROR: Could not rename {file_to_rename}") - error_count += 1 - continue + # Find identified file names and paths + files_to_rename = [re.split(r":", line.strip())[0] for line in lines if re.search(r"FOUND$", line)] + + for file_to_rename in files_to_rename: + new_name = file_to_rename + args.suffix + + if not args.undo: + # Print warning if file can't be found + if not os.path.isfile(file_to_rename): + if os.path.isfile(new_name): + print(f"WARNING: The file {file_to_rename} was already renamed to {new_name}") + warning_count += 1 else: - try: - if not args.dry_run: - os.rename(file_to_rename, new_name) - print(f"SUCCESS: {file_to_rename} -> {new_name}") - success_count += 1 - except Exception as e: - print(f"ERROR: Could not rename {file_to_rename}: {e}") - error_count += 1 + print(f"ERROR: Could not rename {file_to_rename}") + error_count += 1 + continue + # Do the renaming else: - if not os.path.isfile(new_name): - if os.path.isfile(file_to_rename): - print(f"WARNING: The file {file_to_rename} was not yet renamed") - warning_count += 1 - else: - print(f"ERROR: Could not undo renaming of {new_name}") - error_count += 1 - continue - else: - try: - if not args.dry_run: - os.rename(new_name, file_to_rename) - print(f"SUCCESS: {new_name} -> {file_to_rename}") - success_count += 1 - except Exception as e: - print(f"ERROR: Could not undo renaming of {new_name}: {e}") - error_count += 1 + try: + if not args.dry_run: + os.rename(file_to_rename, new_name) + print(f"SUCCESS: {file_to_rename} -> {new_name}") + success_count += 1 + except Exception as e: + print(f"ERROR: Could not rename {file_to_rename}: {e}") + error_count += 1 + # If the changes are to be undone + else: + # Print warning if file can't be found + if not os.path.isfile(new_name): + if os.path.isfile(file_to_rename): + print(f"WARNING: The file {file_to_rename} was not yet renamed") + warning_count += 1 + else: + print(f"ERROR: Could not undo renaming of {new_name}") + error_count += 1 + continue + # Undo the renaming + else: + try: + if not args.dry_run: + os.rename(new_name, file_to_rename) + print(f"SUCCESS: {new_name} -> {file_to_rename}") + success_count += 1 + except Exception as e: + print(f"ERROR: Could not undo renaming of {new_name}: {e}") + error_count += 1 + + # Print result if error_count or warning_count: print(f"--------------------------------------------------------------\n" - f"Renaming {'would have' if args.dry_run else ''} finished with " + f"Renaming {'would have ' if args.dry_run else ''}finished with " f"{error_count} error{'s' if error_count > 1 else ''} and " f"{warning_count} warning{'s' if warning_count > 1 else ''}.\n" - f"{success_count if success_count > 0 else 'No'} files {'would have been ' if args.dry_run else 'were '} " + f"{success_count if success_count > 0 else 'No'} files {'would have been' if args.dry_run else 'were'} " "renamed successfully.") else: print(f"--------------------------------------------------------------\n" - f"{success_count if success_count > 0 else 'No'} files {'would have been ' if args.dry_run else 'were '}" + f"{success_count if success_count > 0 else 'No'} files {'would have been' if args.dry_run else 'were'}" "renamed successfully.")