import argparse

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',
                       '--suffix',
                       type=str,
                       help='suffix to add to the end of the file names',
                       default='.VIRUS')

my_parser.add_argument('-u',
                       '--undo',
                       action='store_true',
                       help='undo renaming')

my_parser.add_argument('--dry-run',
                       action='store_true',
                       help='perform a test run where no file names are actually changed')


def main():
    args = my_parser.parse_args()
    file = args.File
    success_count = 0
    error_count = 0
    warning_count = 0

    if not os.path.isfile(file):
        print(f"The file '{file}' does not exist")
        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)]

        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
                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
            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

    if error_count or warning_count:
        print(f"--------------------------------------------------------------\n"
              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 '} "
              "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 '}"
              "renamed successfully.")


if __name__ == '__main__':
    main()