#! /usr/bin/python3
import argparse
import os
import sys

from bsign_helper import BSignHelper
from hash_pack import ExtSigFlag, ExtSigType, PacketHandler

# Exception type to flags mapping
EXCEPTION_TYPE_TO_FLAGS = {
    # MDWE_CHECK_X_AFTR_W_EXCEPT
    "after-write": (ExtSigType.ANON_HASH, ExtSigFlag.EXCEPTION_1),
    # MDWE_CHECK_X_PLUS_W_EXCEPT
    "plus-write": (ExtSigType.ANON_HASH, ExtSigFlag.EXCEPTION_2),
    # Both flags - complete MDWE bypass
    "mdwe-full": (
        ExtSigType.ANON_HASH,
        ExtSigFlag.EXCEPTION_1 | ExtSigFlag.EXCEPTION_2,
    ),
}


def main():
    parser = argparse.ArgumentParser(
        description="Create special digsig exception signature for binary file",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Available exception types:
  after-write - MDWE exception: binary will be allowed to give EXEC prot to mapping
                even after it had WRITE prot
  plus-write  - MDWE exception: binary will completely skip MDWE mechanism
  mdwe-full   - MDWE restrictions will be ignored for this binary

Examples:
  %(prog)s after-write /usr/bin/ls
  %(prog)s plus-write /usr/bin/ls
        """,
    )

    parser.add_argument(
        "exception_type",
        choices=["after-write", "plus-write", "mdwe-full"],
        help="Exception types:"
        "after-write (MDWE execute after write)"
        "plus-write (MDWE simultaneous execute+write)"
        "mdwe-full (MDWE complete bypass)",
    )

    parser.add_argument("binary_path", help="Path to binary file")

    parser.add_argument("-p", "--passphrase", help="GPG passphrase (for signing)")

    parser.add_argument(
        "-o",
        "--output-dir",
        help="Output directory (default: current directory)",
        default="/etc/digsig/anon_hash/signed/",
    )

    args = parser.parse_args()

    try:
        # Get exception type and flags
        sig_type, sig_flags = EXCEPTION_TYPE_TO_FLAGS[args.exception_type]

        # Get binary hash using bsign
        bsign = BSignHelper(passphrase=args.passphrase)
        real_path = os.path.realpath(args.binary_path)
        hash_data = bsign.get_file_hash_from_detached(real_path)

        if not hash_data:
            print(f"Error: Failed to get hash for {real_path}")
            sys.exit(1)

        print(f"Hash obtained: {hash_data.hex()}")

        base_name = os.path.basename(real_path)
        tmp_path = f"/tmp/{base_name}.except"
        with open(tmp_path, "wb") as f:
            f.write(hash_data + b"\n")

        try:
            # Create output directory if it doesn't exist
            os.makedirs(args.output_dir, exist_ok=True)

            # Sign the hash file
            PacketHandler.sign_file_with_hashes(
                input_file=tmp_path,
                packet_flags=0,
                type=sig_type,
                sig_flags=sig_flags,
                passphrase=args.passphrase,
                output_dir=args.output_dir,
            )

            output_file = (
                os.path.join(args.output_dir, os.path.basename(tmp_path)) + ".sign"
            )
            print(f"Exception created successfully: {output_file}")

        finally:
            # Remove temporary file
            os.unlink(tmp_path)

    except KeyboardInterrupt:
        print("\nOperation interrupted")
        sys.exit(1)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
