#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
----------------------------------------------------------------------------
"MIT License:
ganapati (@G4N4P4T1) is the original author of this project. As long as you retain this notice you
can do whatever you want with this stuff. If we meet some day, and you think
this stuff is worth it, you can buy me a beer in return.
----------------------------------------------------------------------------
"""

import os
import sys
import logging
import argparse
import urllib3
import tempfile
from glob import glob
from os.path import dirname, basename, isfile, join
from urllib3.exceptions import InsecureRequestWarning
from RsaCtfTool.lib.crypto_wrapper import RSA
from RsaCtfTool.lib.rsa_attack import RSAAttack
from RsaCtfTool.lib.number_theory import invmod, factor_ned
from RsaCtfTool.lib.utils import get_numeric_value, print_results, get_base64_value, n2s
from RsaCtfTool.lib.customlogger import CustomFormatter, logger_levels
from RsaCtfTool.lib.keys_wrapper import (
    generate_pq_from_n_and_p_or_q,
    generate_keys_from_p_q_e_n,
    PrivateKey,
)
from RsaCtfTool.lib.idrsa_pub_disector import disect_idrsa_pub
from RsaCtfTool.lib.is_roca_test import is_roca_vulnerable

# Remove insecure warning for factordb.com
urllib3.disable_warnings(InsecureRequestWarning)

# Change recursion limit for... you know, factorizing stuff...
sys.setrecursionlimit(5000)


def banner():
    cEND = "\033[0m"
    cRED = "\033[1;31m"
    return (
        r"""
__________               R_______________________________E __                .__
\______   \ ___________  R\_   ___ \__    ___/\_   _____/E/  |_  ____   ____ |  |
 |       _//  ___/\__  \ R/    \  \/ |    |    |    __)E \   __\/  _ \ /  _ \|  |
 |    |   \\\___ \  / __ \R\     \____|    |    |     \E   |  | (  <_> |  <_> )  |__
 |____|_  /____  >(____  /R\______  /|____|    \___  /E   |__|  \____/ \____/|____/
        \/     \/      \/        R\/E               R\/E

""".replace("R", cRED).replace("E", cEND)
        + """
Disclaimer: this tool is meant for educational purposes, for those doing CTF's first try:

Learning the basis of RSA math, understand number theory, modular arithmetic, integer factorization, fundamental theorem of arithmetic.
Read the code in this repo to see what and how it does and how to improve it, send PR's.
Avoid copy-paste-run and at last run this tool (knowing the math is more valuable than knowing how to run this tool).

"""
    )


def parse_args():
    parser = argparse.ArgumentParser(description="RSA CTF Tool")

    parser.add_argument(
        "--publickey", help="public key file. You can use wildcards for multiple keys."
    )
    parser.add_argument(
        "--output", help="output file for results (private keys, plaintext data)."
    )
    parser.add_argument(
        "--timeout",
        help="Timeout for long attacks in seconds. default is 60s min: MIN_INT in C, max: MAX_INT in C, values < 1 have the same effect as MAX_INT",
        default=60,
        type=int,
    )
    parser.add_argument(
        "--createpub",
        help="Take n and e from cli and just print a public key then exit",
        action="store_true",
    )
    parser.add_argument(
        "--dumpkey",
        help="Just dump the RSA variables from a key - n,e,d,p,q",
        action="store_true",
    )
    parser.add_argument(
        "--ext",
        help="Extended dump of RSA private variables in --dumpkey mode - dp,dq,pinv,qinv).",
        action="store_true",
    )
    parser.add_argument(
        "--decryptfile",
        help="decrypt a file, using commas to separate multiple paths",
        default=None,
    )
    parser.add_argument(
        "--decrypt",
        help="decrypt a cipher, using commas to separate multiple ciphers",
        default=None,
    )
    parser.add_argument(
        "--verbosity", help="verbose mode", choices=logger_levels.keys(), default="INFO"
    )
    parser.add_argument(
        "--private", help="Display private key if recovered", action="store_true"
    )
    parser.add_argument("--tests", help="Run tests on attacks", action="store_true")
    parser.add_argument(
        "--ecmdigits",
        type=int,
        help="Optionally an estimate as to how long one of the primes is for ECM method",
        default=None,
    )
    parser.add_argument("-n", help="Specify the modulus. format : int or 0xhex")
    parser.add_argument(
        "-p", help="Specify the first prime number. format : int or 0xhex"
    )
    parser.add_argument(
        "-q", help="Specify the second prime number. format : int or 0xhex"
    )
    parser.add_argument(
        "-e",
        help="Specify the public exponent, using commas to separate multiple exponents. format : int or 0xhex",
    )
    parser.add_argument(
        "-d",
        help="Specify the private exponent. Format : int or 0xhex",
    )
    parser.add_argument("--key", help="Specify the private key file.")
    parser.add_argument("--password", help="Private key password if needed.")

    # If no arguments, display help and exit
    if len(sys.argv) == 1:
        print(banner())
        parser.print_help()
        sys.exit(1)

    # Dynamic load all attacks for choices in argparse
    attacks = glob(
        join(dirname(os.path.realpath(__file__)), "attacks", "single_key", "*.py")
    )
    attacks += glob(
        join(dirname(os.path.realpath(__file__)), "attacks", "multi_keys", "*.py")
    )

    attacks_filtered = [
        basename(f)[:-3] for f in attacks if isfile(f) and not f.endswith("__init__.py")
    ]
    attacks_list = [_ for _ in attacks_filtered if _ != "nullattack"] + ["all"]
    parser.add_argument(
        "--attack",
        help="Specify the attack modes.",
        default="all",
        nargs="+",
        choices=attacks_list,
    )
    parser.add_argument(
        "--sendtofdb", help="Send results to factordb", action="store_true"
    )
    parser.add_argument(
        "--isconspicuous", help="conspicuous key check", action="store_true"
    )
    parser.add_argument(
        "--isroca", help="Check if given key is roca", action="store_true"
    )

    parser.add_argument(
        "--convert_idrsa_pub", help="Convert idrsa.pub to pem", action="store_true"
    )
    parser.add_argument(
        "--check_publickey",
        help="Check publickey if modulus is well formed before attack",
        action="store_true",
    )

    parser.add_argument(
        "--partial",
        help="work with partial private keys",
        action="store_true",
    )

    parser.add_argument(
        "--cleanup",
        help="cleanup *.pub files after finish",
        action="store_true",
    )

    parser.add_argument(
        "--withtraceback",
        help="show tracebacks",
        action="store_true",
    )

    parser.add_argument(
        "--show_modulus",
        help="show tracebacks",
        action="store_true",
    )

    args = parser.parse_args()
    args.attacks_list = attacks_list
    return args


def run_conspicuous_check(args, logger):
    try:
        pub_key, priv_key = generate_keys_from_p_q_e_n(args.p, args.q, args.e, args.n)
    except ValueError:
        logger.error(
            "Looks like the values for generating key are not ok... (no invmod)"
        )
        return False
    c = priv_key.is_conspicuous()
    if c:
        logger.warning("[!] Key is conspicuous...")
    return c


def run_attacks(args, logger):
    # Run attacks
    attackobj = RSAAttack(args)
    selected_attacks = args.attacks_list

    # Run tests
    if args.publickey is None and args.tests:
        if args.attack is not None:
            if "," not in args.attack:
                selected_attacks = args.attack
        if "all" in selected_attacks:
            selected_attacks = args.attacks_list
        logger.info("Testing attacks: %d" % (len(selected_attacks)-1)) # Exclude "all"

        tmpfile = tempfile.NamedTemporaryFile()
        with open(tmpfile.name, "wb") as tmpfd:
            tmpfd.write(RSA.construct((35, 3)).publickey().exportKey())
            attackobj.attack_single_key(tmpfile.name, selected_attacks, test=True)
            sys.exit(0)

    # Attack multiple keys
    if args.publickey is not None and len(args.publickey) > 1:
        attackobj.attack_multiple_keys(args.publickey, selected_attacks)

    # Attack key
    if args.publickey is not None:
        for publickey in args.publickey:
            attackobj.implemented_attacks = []
            attackobj.decrypted = []
            logger.info("\n[*] Testing key %s." % publickey)
            attackobj.attack_single_key(publickey, selected_attacks)
    if args.publickey is None:
        if args.partial:
            priv_key = PrivateKey(filename=args.key, password=None)
            attackobj.attack_single_key(priv_key, selected_attacks)
        else:
            logger.error("No key specified")
        if args.n is not None:
            # FIXME
            publickey, _privkey = generate_keys_from_p_q_e_n(
                args.p, args.q, args.e, args.n
            )
            attackobj.attack_single_key(publickey, selected_attacks)
    return args


def convert_idrsa_pub(args, logger):
    # for publickey in args.publickey:
    publickeys = glob(args.publickey)
    for publickey in publickeys:
        logger.info(f"Converting {publickey}: to pem...")
        with open(publickey, "r") as key_data_fd:
            for line in key_data_fd:
                n, e = disect_idrsa_pub(line.rstrip())
                if n and e:
                    pub_key, _ = generate_keys_from_p_q_e_n(None, None, e, n)
                    if pub_key:
                        logger.info(pub_key.decode("utf-8"))
                    else:
                        logger.error("Error generating keys from n and e values.")


def check_is_roca(args, logger):
    """
    Checks the given list of public key files for the ROCA vulnerability.

    Args:
        args (Namespace): Command line arguments or configuration settings.
        logger: Logger instance for logging messages.

    Returns:
        bool: True if any of the public keys are vulnerable, False otherwise.
    """
    vuln = False
    pubkeyfilelist = glob(args.publickey)
    for publickey in pubkeyfilelist:
        logger.info(f"[-] Details for {publickey}:")
        with open(publickey, "rb") as key_data_fd:
            try:
                key = RSA.importKey(key_data_fd.read())
            except Exception:
                key = None
                logger.error(f"[!] Error file format: {publickey}")
            if key is not None:
                if is_roca_vulnerable(key.n):
                    vuln = True
                    logger.warning(f"[!] Public key {publickey}: is roca!!!")
                else:
                    logger.info(
                        f"[-] Public key {publickey}: is not roca, you are safe"
                    )
    return vuln


def load_keys(args, logger):
    tmpfile = None
    args.publickey = []
    for e in args.e if isinstance(args.e, list) else [args.e]:
        tmpfile = tempfile.NamedTemporaryFile(delete=False)
        with open(tmpfile.name, "wb") as tmpfd:
            tmpfd.write(
                RSA.construct((args.n, e)).publickey().exportKey(),
            )
        args.publickey.append(tmpfile.name)
    return args


def dump_key_parameters(args):
    key_data = open(args.key, "rb").read()
    key = RSA.importKey(key_data)
    print(f"n: {str(key.n)}")
    print(f"e: {str(key.e)}")
    if key.has_private():
        print(f"d: {str(key.d)}")
        print(f"p: {str(key.p)}")
        print(f"q: {str(key.q)}")
        if args.ext:
            dp = key.d % (key.p - 1)
            dq = key.d % (key.q - 1)
            pinv = invmod(key.p, key.q)
            qinv = invmod(key.q, key.p)
            print(f"dp: {str(dp)}")
            print(f"dq: {str(dq)}")
            print(f"pinv: {str(pinv)}")
            print(f"qinv: {str(qinv)}")


def decrypt_file(args, logger):
    """
    Decrypts files specified in args.decryptfile using the provided private key, or prepares files for decryption with a public key.
    Note: Currently, decryption with a private key may not work correctly.

    Args:
        args (Namespace): Command-line arguments.
        logger (Logger): Logger object for logging messages.

    Returns:
        bool: True if decryption is successful with a private key, or if files are prepared for decryption with a public key; False otherwise.
    """
    decrypt_array = []
    for decrypt in args.decryptfile.split(","):
        try:
            with open(decrypt, "rb") as cipherfile_fd:
                decrypt_value = get_base64_value(cipherfile_fd.read())
                decrypt_array.append(decrypt_value)
        except OSError:
            logger.info("--decryptfile : file not found or not readable.")
            return False
    args.decrypt = decrypt_array

    # Check if a private key is provided and there's something to decrypt
    if args.key and args.decrypt:
        priv_key = PrivateKey(filename=args.key, password=args.password)
        decrypts = priv_key.decrypt(args.decrypt)
        print_results(args, None, priv_key, decrypts)
        return True

    # Check if a public key is provided
    if args.publickey:
        return True

    # Check if n and e are provided
    if args.n and args.e:
        return True

    # No private key or public key provided
    logger.error("Private key or public key and decrypted data are required.")
    return False


def pubkey_detail(args, logger):
    for publickey in args.publickey:
        logger.info(f"Details for {publickey}:")
        with open(publickey, "rb") as key_data_fd:
            key = RSA.importKey(key_data_fd.read())
            print(f"n: {str(key.n)}")
            print(f"e: {str(key.e)}")


def cleanup(args):
    if args.publickey is not None:
        for pub in args.publickey:
            try:
                if "tmp" in pub and "tmp/" not in pub:
                    os.remove(pub)
            except Exception:
                continue


def _setup_logging(args):
    logging.basicConfig(level=logger_levels[args.verbosity])
    ch = logging.StreamHandler(sys.stderr)
    ch.setFormatter(CustomFormatter())
    logger = logging.getLogger("global_logger")
    logger.propagate = False
    logger.addHandler(ch)
    if not args.private and not args.tests:
        logger.warning(
            "private argument is not set, the private key will not be displayed, even if recovered."
        )
    return logger


def _parse_numeric_args(args):
    if args.p is not None:
        args.p = get_numeric_value(args.p)
    if args.q is not None:
        args.q = get_numeric_value(args.q)
    if args.d is not None:
        args.d = get_numeric_value(args.d)
    if args.n is not None:
        args.n = get_numeric_value(args.n)
    if args.e is not None:
        e_array = [get_numeric_value(e) for e in args.e.split(",")]
        args.e = e_array if len(e_array) > 1 else e_array[0]
    elif args.n is not None:
        args.e = 65537
    return args


def _warn_if_extra_primes(args, logger):
    if args.n is not None and (args.p is not None or args.q is not None):
        logger.warning(
            "[!] It seems you already provided one of the prime factors, nothing to do here..."
        )


def _compute_n_from_pq(args):
    if args.n is None and args.p is not None and args.q is not None:
        args.n = args.p * args.q


def _recover_pq_from_ned(args, logger):
    if not (args.n is not None
            and args.e is not None
            and args.d is not None
            and args.p is None
            and args.q is None):
        return
    pq = factor_ned(args.n, args.e, args.d)
    if pq is not None:
        args.p, args.q = pq
    else:
        logger.warning("[!] Impossible to recover p and q from d")


def _compute_other_prime(args):
    if args.n and (args.p or args.q):
        args.p, args.q = generate_pq_from_n_and_p_or_q(args.n, args.p, args.q)


def _compute_missing_values(args, logger):
    _warn_if_extra_primes(args, logger)
    _compute_n_from_pq(args)
    _recover_pq_from_ned(args, logger)
    _compute_other_prime(args)
    return args


def _handle_decrypt_input(args, logger):
    if args.decrypt is not None:
        decrypt_array = []
        for decrypt in args.decrypt.split(","):
            try:
                decrypt = get_numeric_value(decrypt)
            except Exception:
                decrypt = get_base64_value(decrypt)
            decrypt_array.append(n2s(decrypt))
        args.decrypt = decrypt_array
    if args.decryptfile is not None:
        if not decrypt_file(args, logger):
            sys.exit(-1)
    return args


def _handle_early_exit_modes(args, logger):
    if args.convert_idrsa_pub:
        convert_idrsa_pub(args, logger)
        sys.exit(0)
    if args.isroca:
        check_is_roca(args, logger)
        sys.exit(0)
    if args.createpub:
        pub_key, priv_key = generate_keys_from_p_q_e_n(args.p, args.q, args.e, args.n)
        print(pub_key.decode("utf-8"))
        sys.exit(0)
    if (
        args.dumpkey
        and not args.private
        and args.decrypt is None
        and args.decryptfile is None
        and args.publickey is not None
    ):
        pubkey_detail(args, logger)
        sys.exit(0)
    if args.key is not None and args.dumpkey:
        dump_key_parameters(args)
        sys.exit(0)
    if args.key is not None and args.isconspicuous:
        if run_conspicuous_check(args, logger):
            sys.exit(-1)
        sys.exit(0)
    return False


def _load_public_keys(args):
    if args.publickey is None and args.e is not None and args.n is not None:
        args = load_keys(args, logging.getLogger("global_logger"))
    elif args.publickey is not None:
        if "*" in args.publickey or "?" in args.publickey:
            args.publickey = glob(args.publickey)
        elif "," in args.publickey:
            args.publickey = args.publickey.split(",")
        else:
            args.publickey = [args.publickey]
    return args


def _handle_fully_specified_key(args, logger):
    decrypts = []
    if not (
        args.p is not None
        and args.q is not None
        and args.e is not None
        and args.n is not None
    ):
        return None

    try:
        pub_key, priv_key = generate_keys_from_p_q_e_n(args.p, args.q, args.e, args.n)
    except ValueError:
        logger.error(
            "Looks like the values for generating key are not ok... (no invmod)"
        )
        sys.exit(1)

    if args.createpub:
        pub_key, priv_key = generate_keys_from_p_q_e_n(args.p, args.q, args.e, args.n)
        print(pub_key.decode("utf-8"))

    if args.decrypt is not None:
        for u in args.decrypt:
            if priv_key is not None:
                decrypts.append(priv_key.decrypt(args.decrypt))
            else:
                logger.error(
                    "Looks like the values for generating key are not ok... (no invmod)"
                )
                sys.exit(1)
    print_results(args, args.publickey[0], priv_key, decrypts)
    sys.exit(0)
    return decrypts


def main():
    args = parse_args()
    logger = _setup_logging(args)
    args = _parse_numeric_args(args)
    args = _compute_missing_values(args, logger)
    args = _handle_decrypt_input(args, logger)

    if _handle_early_exit_modes(args, logger):
        return

    args = _load_public_keys(args)
    _handle_fully_specified_key(args, logger)
    args = run_attacks(args, logger)

    if args.cleanup:
        cleanup(args)


if __name__ == "__main__":
    main()
