#!/bin/python3
# TODO: Add color highlighting to help

import subprocess
import sys
import os
import colorama
import json

colorama.init(autoreset=True)


# default verbosity state, controlled with `--verbose` flag
VERBOSE = False
# update ability is turned off by default, controlled with
# `--update` flag. This is off to avoid need for user input
# and only work with returned error codes, making it possible
# to run this script easily with cronjobs or similar schedulers
ENABLE_UPDATE = False
# update automatically without asking the user first (will only
# have effect if ENABLE_UPDATE is True - or turned by user), this
# will still inform user about those files being updated, but user
# won't be able to choose not to update
AUTO_UPDATE = False
# default path to the JSON file that stores known file checksums
# this can be overrided by using `--checksum-file=path` flag
CHECKSUM_FILE = '/usr/local/share/tamper-check/checksums.json'
# This shouldn't be touched, entries in this list should be valid
# path strings which are the new files that should be added to
# CHECKSUM_FILE, controlled by flag `--add=/path/to/file`
TO_ADD = []


def _print_help(prepend_newline=False):
    if prepend_newline:
        print()
    print(
        'tamper-check is a command line utility to automate checking for file edits.\n'
        'This is achieved by storing sha256 checksums of each added file and comparing them\n\n'
        'Accepted flags:\n'
        '  `-u`/`--update`: If invalid checksum is found, ask user if it should be updated (y/n)\n'
        '  `-a=path`/`--add=path`: Add a new file to the list of check entries\n'
        '  `--no-confirm`: Used in combination with `--update`, automatically assumes `y` for all questions\n'
        '  `--auto-update`: Combines `--update` and `--no-confirm`\n'
        '  `-f=path/--checksum-file=path: JSON file storing the file checksums, defaults to /usr/local/share/tamper-check/checksums.json\n`'
        '  `-v`/`--verbose`: Verbose mode, show checksums on failures and some more info\n'
        '  `-h`/`--help`: Show this help'
    )


def _yes_no(text, add_yn=True):
    if add_yn:
        text += ' (y/n): '
    while True:
        user_inp = input(text).lower()
        if user_inp in ('y', 'yes'):
            return True
        elif user_inp in ('n', 'no'):
            return False


def _get_checksum_dict(file):
    try:
        with open(file, 'r') as f:
            checksums = json.load(f)
        return checksums
    except FileNotFoundError:
        print(
            f'{colorama.Fore.YELLOW}Checksum file not found: {colorama.Fore.RESET}'
            f"'{colorama.Fore.BLUE}{file}{colorama.Fore.RESET}'{colorama.Fore.YELLOW} "
            'Creating new empty checksum file...'
        )
        os.makedirs(os.path.dirname(file), exist_ok=True)
        checksums = {}
        with open(file, 'w') as f:
            json.dump(checksums, f, indent=4)
        return checksums
    except PermissionError:
        print(
            f'{colorama.Fore.RED}PermissionError: {colorama.Fore.RESET}'
            'to run tamper-check you must have read access to checksum file: '
            f"'{colorama.Fore.BLUE}{file}{colorama.Fore.RESET}' (forgot sudo?)"
        )
        exit(2)



def _add_file(file_path, checksum, checksum_file, update_only=False):
    checksums = _get_checksum_dict(checksum_file)

    if file_path in checksums and not update_only:
        print(
            f"{colorama.Fore.RED}Path {colorama.Fore.RESET}"
            f"'{colorama.Fore.BLUE}{file_path}{colorama.Fore.RESET}' {colorama.Fore.RED}"
            "is already in the checksum file perhaps you wanted `--update`?"
        )
        return False

    checksums[file_path] = checksum
    try:
        with open(checksum_file, 'w') as f:
            json.dump(checksums, f, indent=4)
    except PermissionError:
        print(
            f'{colorama.Fore.RED}PermissionError: {colorama.Fore.RESET}'
            'To add a new rule, you must have write access to: '
            f"'{colorama.Fore.BLUE}{checksum_file}{colorama.Fore.RESET}' (forgot sudo?)"
        )
        exit(2)

    return True


def _get_checksum(file):
    proc = subprocess.run(['sha256sum', file], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    proc_stdout = proc.stdout.decode('utf-8')
    if "No such file or directory" in proc_stdout:
        print(
            f'{colorama.Fore.RED}FileNotFound: {colorama.Fore.RESET}'
            f"'{colorama.Fore.BLUE}{file}{colorama.Fore.RESET}' not found, can't produce "
            "sha256 checksum, check the checksum file and remove this entry"
        )
        exit(2)
    elif "Permission denied" in proc_stdout:
        raise PermissionError(f"PermissionError: Unable to read file '{file}'")
    return proc_stdout.replace(f'  {file}\n', '')


def _command_exists(command):
    proc = subprocess.run(f'which {command}', stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True)
    return proc.returncode == 0


def ask_update(file_path, checksum_file):
    new_checksum = _get_checksum(file_path)

    file_line = f"'{colorama.Fore.BLUE}{file_path}{colorama.Fore.RESET}'"
    if AUTO_UPDATE or _yes_no('  - ' + file_line + ' update checksum?'):
        if AUTO_UPDATE:
            print('  - updating checksum for ' + file_line)  # path won't get printed by _yes_no(), so print here
        result = _add_file(file_path, new_checksum, checksum_file, update_only=True)
        print(f'{colorama.Fore.GREEN}  -> Updated')
        return result

    print(f'{colorama.Fore.RED}  -> Staying mismatched')
    return False


def run_check():
    checksums = _get_checksum_dict(CHECKSUM_FILE)
    not_matched = []
    for file, checksum in checksums.items():
        try:
            sha256_sum = _get_checksum(file)
        except PermissionError:
            print(
                f"Checksum of '{colorama.Fore.BLUE}{file}{colorama.Fore.RESET}': "
                f'{colorama.Fore.YELLOW}SKIPPED [PermissionError - no read perms]')
            continue

        if sha256_sum != checksum:
            not_matched.append(file)
            print(
                f"Checksum of '{colorama.Fore.BLUE}{file}{colorama.Fore.RESET}': "
                f'{colorama.Fore.RED}FAIL [Checksum Mismatch]'
            )
            if VERBOSE:
                print(f"   -> detected: {colorama.Fore.CYAN}{sha256_sum}")
                print(f"   -> stored:   {colorama.Fore.CYAN}{checksum}")
        else:
            print(f"Checksum of '{colorama.Fore.BLUE}{file}{colorama.Fore.RESET}': {colorama.Fore.GREEN}OK")
    return not_matched


def analyze_args():
    # Using globals isn't usually ideal solution,
    # but it is the easiest way to handle this
    global VERBOSE, ENABLE_UPDATE, AUTO_UPDATE, TO_ADD, CHECKSUM_FILE

    try:
        args = sys.argv[1:]
    except IndexError:
        return

    for arg in args:
        if arg in ('-u', '--update'):
            ENABLE_UPDATE = True
        elif arg == '--no-confirm':
            AUTO_UPDATE = True
        elif arg == '--auto-update':
            ENABLE_UPDATE = True
            AUTO_UPDATE = True
        elif arg.startswith('--add=') or arg.startswith('-a='):
            path = arg.replace('--add=', '').replace('-a=', '')
            if os.path.exists(path):
                path = os.path.abspath(path)
                TO_ADD.append(path)
            else:
                print(
                    f'{colorama.Fore.RED}FileNotFoundError: {colorama.Fore.RESET}'
                    f"'{colorama.Fore.BLUE}{path}{colorama.Fore.RESET}' -> invalid path"
                )
                exit(2)
        elif arg in ('-v', '--verbose'):
            VERBOSE = True
        elif arg.startswith('-f=') or arg.startswith('--checksum-file='):
            path = arg.replace('--checksum-file=', '').replace('-f=', '')
            if os.path.exists(path):
                path = os.path.abspath(path)
                CHECKSUM_FILE = path
            else:
                print(
                    f'{colorama.Fore.RED}FileNotFoundError: {colorama.Fore.RESET}'
                    f"'{colorama.Fore.BLUE}{path}{colorama.Fore.RESET}' -> invalid path"
                )
        elif arg in ('-h', '--help'):
            _print_help()
            exit()
        else:
            print(f'{colorama.Fore.RED}Unrecognized flag: {colorama.Fore.YELLOW}{arg}')
            _print_help(prepend_newline=True)
            exit(2)


if __name__ == '__main__':
    analyze_args()

    # if file adding was requested, handle it
    if len(TO_ADD) > 0:
        for file in TO_ADD:
            checksum = _get_checksum(file)
            if _add_file(file, checksum, CHECKSUM_FILE):
                print(
                    f"Added '{colorama.Fore.YELLOW}{file}{colorama.Fore.RESET}': "
                    f"'{colorama.Fore.CYAN}{checksum}{colorama.Fore.RESET}'"
                )
        exit(0)  # when adding files, don't run the check too

    # run the actual checksum verifier
    not_matched = run_check()

    if len(not_matched) > 0 and ENABLE_UPDATE is True:
        print('\nFiles with mismatched checksums:')
        unfixed = not_matched.copy()
        for mismatched_file in not_matched:
            if ask_update(mismatched_file, CHECKSUM_FILE) is True:
                # User did not choose to update the checksum
                unfixed.remove(mismatched_file)
        if len(unfixed) > 0:
            exit(1)  # we still have some unfixed mismatched checksums, exit with 1

    elif len(not_matched) > 0:
        print('\nFiles with mismatched checksums:')
        for mismatched_file in not_matched:
            print(
                f'{colorama.Fore.RED}  - '
                f"{colorama.Fore.RESET}'{colorama.Fore.BLUE}{mismatched_file}{colorama.Fore.RESET}'"
            )
        exit(1)  # exit with error code in case we have changed checksums
    else:
        print(f'\n{colorama.Fore.GREEN}All checksums are correct')