shacheck.py

"""
Script to check the hash of a file
and compare it to the reference hash.
Only works with SHA algorithms.

usage: python shacheck.py [-h] [--abs] algorithm filename ref_hash

positional arguments:
  algorithm   the SHA algorithm used for the hash generation (e.g. 256)
  filename    relative path of the file to check the ref_hash with
  ref_hash    hashsum to compare to the hash of the file

options:
  -h, --help  show this help message and exit
  --abs       Using an absolute filepath
"""

import argparse
import subprocess
import platform
import sys
from os import path


COLORS = {
    "red": "\x1B[31m",
    "green": "\x1B[32m",
    "reset": "\x1B[0m"
}


def get_args() -> argparse.Namespace:
    """Parse command line args into list"""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "algorithm",
        help="the SHA algorithm used for the hash generation (e.g. 256)"
    )
    parser.add_argument(
        "--abs",
        help="Using an absolute filepath",
        action="store_true",
        required=False
    )
    parser.add_argument(
        "filename",
        help="relative path of the file to check the ref_hash with"
    )
    parser.add_argument(
        "ref_hash",
        help="hashsum to compare to the hash of the file"
    )
    return parser.parse_args()


def compare(hash: str, ref: str) -> None:
    """
    Comparing the file hash with the reference hash
    displaying the result to the console.
    """
    if hash == ref:
        print(f"{COLORS['green']}success{COLORS['reset']}")
    else:
        print(
            f"{COLORS['red']}failure{COLORS['reset']}",
            f"The hash of the file is '{hash}'",
            sep="\n"
            )


def is_win() -> bool:
    """Returns True if platform is Windows"""
    return True if platform.system() == "Windows" else False


def construct_filepath(args: argparse.Namespace) -> str:
    """Constructs the filepath"""
    if args.abs:
        return path.abspath(args.filename)
    else:
        return path.join(path.dirname(path.realpath(__file__)), args.filename)


def main() -> None:
    """main script execution"""
    args = get_args()
    filepath = construct_filepath(args)
    if is_win():
        # Windows system
        with subprocess.Popen(
            ["certutil", "-hashfile", filepath, f"SHA{args.algorithm}"],
            stdout=subprocess.PIPE
        ) as proc:
            res = proc.communicate()[0]
        if proc.returncode == 0:
            shasum = res.decode("utf-8").split("\n")[1].strip()
    else:
        # *nix system (macos, linux)
        with subprocess.Popen(
            ["shasum", "-a", args.algorithm, filepath],
            stdout=subprocess.PIPE
        ) as proc:
            res = proc.communicate()[0]
        if proc.returncode == 0:
            shasum = res.decode("utf-8").split(" ")[0]
    compare(shasum, args.ref_hash)
    print()


if __name__ == "__main__":
    main()

TODO

Make commands secure (escape illegal commands/chars)