#!/usr/bin/env python
import argparse
import os
import re
import sys

# 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 be appended to the end of the filenames',
                       default='.VIRUS')
argparser.add_argument('-u',
                       '--undo',
                       action='store_true',
                       help='Undo renaming')
argparser.add_argument('-d',
                       '--dry-run',
                       action='store_true',
                       help='Perform a test run where no file names are actually changed')


def main():
    args = argparser.parse_args()
    success_count = 0
    error_count = 0
    warning_count = 0

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

    # Read input file
    with args.file as file:
        lines = file.readlines()

    # 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:
                    print(f"ERROR: Could not rename {file_to_rename}")
                    error_count += 1
                continue
            # Do the renaming
            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

        # 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"{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()