import hashlib
import warnings

import cryptography
from cryptography import x509
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import (dh, dsa, ec, ed448, ed25519, mldsa, mlkem,
                                                       rsa, x448, x25519)

from .allkeys import blocklist
from .dsakeys import dsasparse
from .rsakeys import (fermat, pattern, roca, rsabias, rsainvalid, rsapoly, rsarecover, rsawarnings,
                      sharedprimes, smalld, smallfactors, xzbackdoor)

# List of available checks
defaultchecks = {
    "fermat": {
        "type": "rsa",
        "function": fermat,
        "desc": "Fermat factorization / close prime vulnerability",
    },
    "roca": {
        "type": "rsa",
        "function": roca,
        "desc": "Return of the Coopersmith Attack (ROCA) vulnerability",
    },
    "rsainvalid": {
        "type": "rsa",
        "function": rsainvalid,
        "desc": "RSA keys with invalid values",
    },
    "rsapoly": {
        "type": "rsa",
        "function": rsapoly,
        "desc": "0-byte pattern in modulus allowing polynomial factoring",
    },
    "sharedprimes": {
        "type": "rsa",
        "function": sharedprimes,
        "desc": "Shared prime factors (batchgcd)",
    },
    "smalld": {
        "type": "rsa",
        "function": smalld,
        "desc": "Small private d (Wiener's attack)",
    },
    "smallfactors": {
        "type": "rsa",
        "function": smallfactors,
        "desc": "Small prime factors (<=65537, usually corrupt)",
    },
    "blocklist": {
        "type": "all",
        "function": blocklist,
        "desc": "Blocklists of compromised keys",
    },
    "dsasparse": {
        "type": "dsa",
        "function": dsasparse,
        "desc": "DSA sparse key vulnerability"
    }
}

warningchecks = {
    "rsawarnings": {
        "type": "rsa",
        "function": rsawarnings,
        "desc": "RSA key size and exponent warnings",
    },
    "rsabias": {
        "type": "rsa",
        "function": rsabias,
        "desc": "RSA modulus with bias of 0/1 bits",
    },
}

# These lead to a problematic rate of false positives and
# have limited value. They may be removed in the future.
extrachecks = {
    "pattern": {
        "type": "rsa",
        "function": pattern,
        "desc": "Implausible repetition pattern in modulus",
    },
    "xzbackdoor": {
        "type": "rsa",
        "function": xzbackdoor,
        "desc": "xz backdoor payload in RSA n",
    },
}

allchecks = defaultchecks | warningchecks | extrachecks

# cryptography warns about SSH DSA keys being deprecated.
# For now, disable the warnings. Needs a better long-term solution.
warnings.filterwarnings("ignore", category=cryptography.utils.CryptographyDeprecationWarning,
                        module="badkeys.checks")


def _checkkey(key, checks, keyrecover=False):
    r = {}
    if isinstance(key, rsa.RSAPublicKey):
        r["type"] = "rsa"
        r["n"] = key.public_numbers().n
        r["e"] = key.public_numbers().e
        r["bits"] = r["n"].bit_length()
        r["results"] = checkrsa(r["n"], e=r["e"], checks=checks, keyrecover=keyrecover)
    elif isinstance(key, ec.EllipticCurvePublicKey):
        r["type"] = "ec"
        r["curve"] = key.curve.name
        # shorter names for standard curves
        if r["curve"].startswith("sec") and r["curve"].endswith("r1"):
            r["curve"] = r["curve"][3:-2]
        r["x"] = key.public_numbers().x
        r["y"] = key.public_numbers().y
        r["results"] = checkall(r["x"], checks=checks)
    elif isinstance(key, dsa.DSAPublicKey):
        r["type"] = "dsa"
        r["y"] = key.public_numbers().y
        r["p"] = key.parameters().parameter_numbers().p
        r["q"] = key.parameters().parameter_numbers().q
        r["g"] = key.parameters().parameter_numbers().g
        r["bits"] = key.key_size
        # we currently have no checks using the q, g parameters,
        # they could be added in the future
        r["results"] = checkdsa(r["y"], p=r["p"], checks=checks)
    elif isinstance(key, (ed25519.Ed25519PublicKey, x25519.X25519PublicKey,
                          x448.X448PublicKey, ed448.Ed448PublicKey)):
        r["type"] = "ec"
        r["curve"] = str(type(key).__name__).lower()[:-9]
        # convert the raw key into an integer for the blocklist check
        pub_b = key.public_bytes(encoding=serialization.Encoding.Raw,
                                 format=serialization.PublicFormat.Raw)
        r["pub"] = int.from_bytes(pub_b, byteorder="big")
        r["results"] = checkall(r["pub"], checks=checks)
    elif isinstance(key, dh.DHPublicKey):
        r["type"] = "dh"
        try:
            r["y"] = key.public_numbers().y
        except ValueError:
            # happens with e.g. very small (<512) DH keys
            return {"type": "unparseable", "results": {}}
        r["results"] = checkall(r["y"], checks=checks)
    elif isinstance(key, (mldsa.MLDSA44PublicKey, mldsa.MLDSA65PublicKey, mldsa.MLDSA87PublicKey,
                          mlkem.MLKEM768PublicKey, mlkem.MLKEM1024PublicKey)):
        r["type"] = type(key).__name__[0:-9].lower()
        r["pub"] = key.public_bytes_raw()
        r["results"] = checkall(r["pub"], checks=checks)
    else:
        r["type"] = "unsupported"
        r["results"] = {}
    spki = key.public_bytes(serialization.Encoding.DER,
                            serialization.PublicFormat.SubjectPublicKeyInfo)
    r["spkisha256"] = hashlib.sha256(spki).hexdigest()
    return r


def checkrsa(n, e=65537, checks=defaultchecks.keys(), keyrecover=False):
    results = {}
    for check in checks:
        callcheck = allchecks[check]["function"]
        if allchecks[check]["type"] == "rsa":
            r = callcheck(n, e=e)
        elif allchecks[check]["type"] == "all":
            r = callcheck(n)
        else:
            continue
        if r is not False:
            pemkey = None
            if keyrecover and "p" in r and "q" in r:
                pemkey = rsarecover(p=r["p"], q=r["q"], n=n, e=e)
            elif keyrecover and "d" in r:
                pemkey = rsarecover(d=r["d"], n=n, e=e)
            if pemkey:
                r["privatekey"] = pemkey
            results[check] = r
    return results


def checkdsa(y, p=0, checks=defaultchecks.keys()):
    results = {}
    for check in checks:
        callcheck = allchecks[check]["function"]
        if allchecks[check]["type"] == "dsa":
            r = callcheck(y, p=p)
        elif allchecks[check]["type"] == "all":
            r = callcheck(y)
        else:
            continue
        if r is not False:
            results[check] = r
    return results


def checkall(x, checks=defaultchecks.keys()):
    results = {}
    for check in checks:
        if allchecks[check]["type"] != "all":
            continue
        callcheck = allchecks[check]["function"]
        r = callcheck(x)
        if r is not False:
            results[check] = r
    return results


def _reterr(rtype, ex):
    return {"type": rtype, "results": {}, "exception": ex.__class__.__name__, "errmsg": str(ex)}


def checkpubkey(rawkey, checks=defaultchecks.keys(), keyrecover=False):
    try:
        key = serialization.load_pem_public_key(rawkey.encode())
    except ValueError as e:
        # ValueError: defect inputs (partial keys, ASN.1 errors), ECDSA with explicit curves,
        #             unknown key types (e.g., ML-DSA)
        return _reterr("unparseable", e)
    except cryptography.exceptions.UnsupportedAlgorithm as e:
        # UnsupportedAlgorithm: Unsupported elliptic curves
        return _reterr("unsupported", e)
    return _checkkey(key, checks, keyrecover=keyrecover)


def checkprivkey(rawkey, checks=defaultchecks.keys(), keyrecover=False):
    try:
        priv = serialization.load_pem_private_key(rawkey.encode(), password=None)
    except ValueError as e:
        # ValueError: defect inputs, unknown key types, explicit curves
        return _reterr("unparseable", e)
    except (cryptography.exceptions.UnsupportedAlgorithm, TypeError) as e:
        # UnsupportedAlgorithm: Unsupported elliptic curves
        # TypeError: passwords-protected keys
        return _reterr("unsupported", e)
    return _checkkey(priv.public_key(), checks, keyrecover=keyrecover)


def checkcrt(rawcert, checks=defaultchecks.keys(), keyrecover=False):
    try:
        crt = x509.load_pem_x509_certificate(rawcert.encode())
    except (ValueError, cryptography.x509.base.InvalidVersion) as e:
        return _reterr("unparseable", e)
    try:
        pubkey = crt.public_key()
    except (cryptography.exceptions.UnsupportedAlgorithm, ValueError, NotImplementedError) as e:
        # ValueError: unknown key types, explicit curves
        # UnsupportedAlgorithm: unsupported curves
        # NotImplementedError: ? (possibly certificate extension issues)
        return _reterr("unsupported", e)
    return _checkkey(pubkey, checks, keyrecover=keyrecover)


def checkcsr(rawcsr, checks=defaultchecks.keys(), keyrecover=False):
    try:
        csr = x509.load_pem_x509_csr(rawcsr.encode())
    except (ValueError, cryptography.x509.base.InvalidVersion) as e:
        return _reterr("unparseable", e)
    try:
        pubkey = csr.public_key()
    except (cryptography.exceptions.UnsupportedAlgorithm, ValueError) as e:
        # ValueError: unknown key types, explicit curves
        # UnsupportedAlgorithm: unsupported curves
        return _reterr("unsupported", e)
    return _checkkey(pubkey, checks, keyrecover=keyrecover)


def checksshprivkey(sshkey, checks=defaultchecks.keys(), keyrecover=False):
    try:
        pkey = serialization.load_ssh_private_key(sshkey.encode(), password=None)
    except ValueError as e:
        # ValueError: defect keys
        return _reterr("unparseable", e)
    except (cryptography.exceptions.UnsupportedAlgorithm, TypeError) as e:
        # UnsupportedAlgorithm: unsupported key types, e.g., sk-*
        # TypeError: password-protected keys
        return _reterr("unsupported", e)
    return _checkkey(pkey.public_key(), checks, keyrecover=keyrecover)


def checksshpubkey(sshkey, checks=defaultchecks.keys(), keyrecover=False):
    try:
        pkey = serialization.load_ssh_public_key(sshkey.encode())
    except ValueError as e:
        # ValueError: defect keys, non-standard DSA keys (!=1024 bit)
        return _reterr("unparseable", e)
    except (cryptography.exceptions.UnsupportedAlgorithm, NotImplementedError) as e:
        # UnsupportedAlgorithm: unsupported key types
        # NotImplementedError: unsupported EC point format
        return _reterr("unsupported", e)
    return _checkkey(pkey, checks, keyrecover=keyrecover)


def detectandcheck(inkey, checks=defaultchecks.keys(), keyrecover=False):
    if "-----BEGIN CERTIFICATE-----" in inkey:
        return checkcrt(inkey, checks, keyrecover=keyrecover)
    if "-----BEGIN CERTIFICATE REQUEST-----" in inkey:
        return checkcsr(inkey, checks, keyrecover=keyrecover)
    if "-----BEGIN PUBLIC KEY-----" in inkey:
        return checkpubkey(inkey, checks, keyrecover=keyrecover)
    if "-----BEGIN RSA PUBLIC KEY-----" in inkey:
        return checkpubkey(inkey, checks, keyrecover=keyrecover)
    if "-----BEGIN PRIVATE KEY-----" in inkey:
        return checkprivkey(inkey, checks, keyrecover=keyrecover)
    if "-----BEGIN RSA PRIVATE KEY-----" in inkey:
        return checkprivkey(inkey, checks, keyrecover=keyrecover)
    if "-----BEGIN DSA PRIVATE KEY-----" in inkey:
        return checkprivkey(inkey, checks, keyrecover=keyrecover)
    if "-----BEGIN EC PRIVATE KEY-----" in inkey:
        return checkprivkey(inkey, checks, keyrecover=keyrecover)
    if "-----BEGIN OPENSSH PRIVATE KEY-----" in inkey:
        return checksshprivkey(inkey, checks, keyrecover=keyrecover)
    if inkey.startswith(("ssh-", "ecdsa-")):
        return checksshpubkey(inkey, checks, keyrecover=keyrecover)
    return {"type": "notfound", "results": {}}
