import argparse
import json
import re
import signal
import sys

from . import __version__
from .allkeys import loadextrabl, urllookup
from .checks import (_checkkey, allchecks, checkcrt, checkrsa, checksshpubkey, defaultchecks,
                     detectandcheck, warningchecks)
from .dkim import parsedkim
from .dnssec import checkdnskey
from .jwk import checkjwk
from .scanssh import scanssh
from .scantls import scantls
from .update import update_bl
from .utils import _errexit, _esc, _getret, _setret, _warnmsg

MAXINPUTSIZE = 2048000

count = 0

PRECRT = "-----BEGIN CERTIFICATE-----\n"
POSTCRT = "\n-----END CERTIFICATE-----\n"

_parseerrmsg = {
    "privleak": "Private instead of public key",
    "keyparseerror": "Error decoding DER/ASN.1 key structure",
    "expectedrsa": "Wrong key type, expected RSA key",
    "wrongkeylength": "Wrong key length",
    "unknowntype": "Unknown key type",
}


def _sighandler(_signum, _handler):
    print(f"{count} keys processed", file=sys.stderr)


def _printresults(key, where, args):
    if args.json:
        jout = key
        jout["where"] = where
        for val in ["n", "e", "x", "y"]:
            if val in jout:
                jout[val] = f"{jout[val]:x}"
        for result in jout["results"]:
            for val in ["p", "q", "e"]:
                if val in jout["results"][result]:
                    jout["results"][result][val] = f'{jout["results"][result][val]:x}'
        if args.all or jout["results"] or \
           key["type"] in ["unsupported", "unparseable", "notfound"]:
            print(json.dumps(jout))
        return
    kn = key["type"]
    if "bits" in key:
        kn += f"[{key['bits']}]"
    if "curve" in key:
        kn += f"[{key['curve']}]"
    if key["type"] == "unsupported":
        _warnmsg(f"Unsupported key type, {_esc(where)}")
    elif key["type"] == "unparseable":
        reason = ""
        if reason in key and key["reason"] in _parseerrmsg:
            reason = f" ({_parseerrmsg[key['reason']]})"
        elif "reason" in key:
            reason = f" ({key['reason']})"
        _warnmsg(f"Unparseable input{reason}, {_esc(where)}")
    elif key["type"] == "notfound":
        _warnmsg(f"No key found, {_esc(where)}")
    elif args.verbose or args.all:
        if key["results"] == {}:
            print(f"{kn} key ok, {_esc(where)}")
    for check, result in key["results"].items():
        sub = ""
        if "subtest" in result:
            sub = f"/{result['subtest']}"
        print(f"{check}{sub} vulnerability, {kn}, {_esc(where)}")
        _setret(4)
        if args.url and "lookup" in result:
            url, _ = urllookup(result["blid"], result["lookup"])
            if url:
                print(url)
            else:
                _warnmsg("URL lookup failed, not found")
        if args.verbose and "debug" in result:
            print(result["debug"])
        if args.verbose and "p" in result:
            print(f"RSA p {result['p']:02x}")
        if args.verbose and "q" in result:
            print(f"RSA q {result['q']:02x}")
        if args.keyrecover and "privatekey" in result:
            print(result["privatekey"], end="")


def runcli():
    global count
    try:
        signal.signal(signal.SIGHUP, _sighandler)
    except AttributeError:  # for OSes without SIGHUP (e.g. Windows)
        pass

    ap = argparse.ArgumentParser()
    ap.add_argument("infiles", nargs="*", help="Input file (certificate, csr or public key)")
    ap.add_argument("-c", "--checks", help="Comma-separated list of checks (default: all)")
    ap.add_argument("--list", action="store_true", help="Show list of possible checks")
    ap.add_argument("-m", "--moduli", action="store_true",
                    help="Input file is list of RSA hex moduli")
    ap.add_argument("--crt-lines", action="store_true", help="Input file is list of base64 certs")
    ap.add_argument("--ssh-lines", action="store_true",
                    help="Input file is list of ssh public keys")
    ap.add_argument("--dkim", action="store_true", help="Scan DKIM records (in files)")
    ap.add_argument("--dkim-dns", action="store_true",
                    help="Scan DKIM DNS record (hostnames instead of files)")
    ap.add_argument("--dnssec", action="store_true", help="Scan DNSKEY/DNSSEC records (in files)")
    ap.add_argument("--jwk", action="store_true", help="Scan JSON Web Keys / Key Sets")
    ap.add_argument("-a", "--all", action="store_true", help="Show all keys")
    ap.add_argument("-w", "--warnings", action="store_true",
                    help="Enable extra warnings (keysize etc.)")
    ap.add_argument("-v", "--verbose", action="store_true", help="Verbose output")
    ap.add_argument("-j", "--json", action="store_true", help="JSON output")
    ap.add_argument("-q", "--quiet", action="store_true", help="Quiet output (for update commands)")
    ap.add_argument("-u", "--url", action="store_true", help="Show private key URL if possible")
    ap.add_argument("--update-bl", action="store_true", help="Update blocklist")
    ap.add_argument("-U", "--update-bl-and-urls", action="store_true",
                    help="Update blocklist and optional URL lookup list")
    ap.add_argument("--extrabl", help="comma-separated list of extra blocklist files")
    ap.add_argument("-t", "--tls", action="store_true",
                    help="Scan TLS (pass hostnames or IPs instead of files)")
    # default ports for https, smtps, imaps, pop3s, ldaps, ftps
    # and 8443 as a common non-default https port
    ap.add_argument("--tls-ports", default="443,465,636,990,993,995,8443",
                    help="TLS ports (comma-separated)")
    ap.add_argument("-s", "--ssh", action="store_true",
                    help="Scan SSH (pass hostnames or IPs instead of files)")
    ap.add_argument("--ssh-ports", default="22", help="SSH ports (comma-separated)")
    ap.add_argument("--keyrecover", action="store_true",
                    help="If possible, recover and output private key")
    ap.add_argument("--version", action="version", version=__version__)
    args = ap.parse_args()

    if (args.moduli and args.crt_lines) or (args.moduli and args.ssh_lines) \
       or (args.ssh_lines and args.crt_lines):
        _errexit("Multiple input format parameters cannot be combined.")

    if (args.moduli or args.crt_lines or args.ssh_lines) and \
       (args.tls or args.ssh or args.dkim_dns):
        _errexit("Scan modes and input file modes cannot be combined.")

    if args.update_bl_and_urls:
        update_bl(lookup=True, quiet=args.quiet)
        sys.exit()
    if args.update_bl:
        update_bl(quiet=args.quiet)
        sys.exit()

    if args.list:
        for k, v in allchecks.items():
            print(f"{k}/{v['type']} keys: {v['desc']}")
        sys.exit()

    if not args.infiles:
        ap.print_help()
        sys.exit(1)

    if args.extrabl:
        extrabl = args.extrabl.split(",")
        for bl in extrabl:
            loadextrabl(bl)

    if args.checks:
        userchecks = set()
        for c in args.checks.split(","):
            if c == "default":
                userchecks |= defaultchecks.keys()
            elif c == "warning":
                userchecks |= warningchecks.keys()
            elif c in allchecks:
                userchecks.add(c)
            else:
                _errexit(f"{c} is not a valid check")
    elif args.warnings:
        userchecks = defaultchecks.keys() | warningchecks.keys()
    else:
        userchecks = defaultchecks.keys()

    if args.tls:
        ports = [int(p) for p in args.tls_ports.split(",")]
        for host in args.infiles:
            for port in ports:
                keys = scantls(host, port, userchecks)
                for k in keys:
                    _printresults(k, f"tls:{host}:{port}", args)

    if args.ssh:
        ports = [int(p) for p in args.ssh_ports.split(",")]
        for host in args.infiles:
            for port in ports:
                keys = scanssh(host, port)
                for k in keys:
                    _printresults(k, f"ssh:{host}:{port}", args)

    if args.dkim_dns:
        try:
            import dns.resolver
        except ModuleNotFoundError:
            _errexit("DKIM DNS record scanning needs dnspython")
        for host in args.infiles:
            try:
                records = dns.resolver.resolve(host, "TXT").response
            except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
                _warnmsg(f"No TXT record found, {_esc(host)}")
                continue
            found = False
            for record in records.answer[-1]:
                dk = b"".join(record.strings).decode()
                key = parsedkim(dk)

                if not key:
                    continue
                found = True
                if isinstance(key, str):
                    r = {"type": "unparseable", "reason": key, "results": {}}
                else:
                    r = _checkkey(key, checks=userchecks, keyrecover=args.keyrecover)
                _printresults(r, host, args)
            if not found:
                _warnmsg(f"No DKIM/DomainKeys key in TXT record, {_esc(host)}")

    if args.jwk:
        for fn in args.infiles:
            with open(fn) as f:
                try:
                    j = json.load(f)
                except (json.decoder.JSONDecodeError, UnicodeDecodeError, ValueError):
                    _warnmsg(f"No valid JSON, {_esc(fn)}")
                    continue
            if isinstance(j, dict) and "kty" in j:
                r = checkjwk(j, checks=userchecks, keyrecover=args.keyrecover)
                _printresults(r, fn, args)
            elif isinstance(j, dict) and "keys" in j:
                for k in j["keys"]:
                    r = checkjwk(k, checks=userchecks, keyrecover=args.keyrecover)
                    _printresults(r, fn, args)
            else:
                _warnmsg(f"No JWK/JWKS, {_esc(fn)}")

    if args.ssh or args.tls or args.dkim_dns or args.jwk:
        sys.exit(_getret())

    if args.dnssec:
        dnskeyre = re.compile(r"[0-9]{1,3}\s+[0-9]{1,3}\s+[0-9]{1,3}\s+[A-Za-z0-9/+= ]+")

    for fn in args.infiles:
        if fn == "-":
            f = sys.stdin
        else:
            f = open(fn, errors="replace")

        if args.moduli or args.ssh_lines or args.crt_lines or args.dkim:
            lcount = 0

        if args.moduli:
            for line in f:
                try:
                    if line.startswith("Modulus="):
                        n = int(line[8:], 16)
                    else:
                        n = int(line, 16)
                except ValueError:
                    r = {"type": "unparseable", "results": {}}
                    _printresults(r, f"line {count}", args)
                    continue
                if n <= 0:
                    r = {"type": "unparseable", "results": {}}
                    _printresults(r, f"line {count}", args)
                    continue
                r = {"type": "rsa", "bits": n.bit_length()}
                r["results"] = checkrsa(n, checks=userchecks, keyrecover=args.keyrecover)
                _printresults(r, f"modulus {n:02x}", args)
                count += 1
                lcount += 1
        elif args.crt_lines:
            lno = 0
            for line in f:
                lno += 1
                desc = f"{fn}[{lno}]"
                ll = re.split("[,; ]", line.rstrip(), maxsplit=1)
                if len(ll) == 2:
                    desc += f" {ll[1]}"
                crt = PRECRT + ll[0] + POSTCRT
                r = checkcrt(crt, checks=userchecks, keyrecover=args.keyrecover)
                _printresults(r, desc, args)
                count += 1
                lcount += 1
        elif args.ssh_lines:
            lno = 0
            for sline in f:
                lno += 1
                line = sline.strip()
                if line == "" or line.startswith("#"):
                    continue
                desc = f"{fn}[{lno}]"
                if not line.startswith(("ssh-", "ecdsa-")):
                    if "ssh-" in line:
                        line = "ssh-" + line.split("ssh-", 1)[1]
                    elif "ecdsa-" in line:
                        line = "ecdsa-" + line.split("ecdsa-", 1)[1]
                    else:
                        _warnmsg(f"No SSH key in {desc}")
                        continue
                ll = line.split(" ", 2)
                if len(ll) == 3:
                    desc += f" {ll[2]}"
                r = checksshpubkey(line, checks=userchecks, keyrecover=args.keyrecover)
                _printresults(r, desc, args)
                count += 1
                lcount += 1
        elif args.dkim:
            lno = 0
            for line in f:
                lno += 1
                desc = f"{fn}[{lno}]"
                key = parsedkim(line)
                if not key:
                    continue
                if isinstance(key, str):
                    r = {"type": "unparseable", "reason": key, "results": {}}
                else:
                    r = _checkkey(key, checks=userchecks, keyrecover=args.keyrecover)
                _printresults(r, desc, args)
                count += 1
                lcount += 1
        elif args.dnssec:
            fcontent = f.read(MAXINPUTSIZE)

            keyrecs = dnskeyre.findall(fcontent)
            if not keyrecs:
                _warnmsg(f"No DNSSEC key found, {_esc(fn)}")
            for rec in keyrecs:
                r = checkdnskey(rec, checks=userchecks, keyrecover=args.keyrecover)
                _printresults(r, fn, args)
                count += 1
        else:
            fcontent = f.read(MAXINPUTSIZE)
            r = detectandcheck(fcontent, checks=userchecks, keyrecover=args.keyrecover)
            _printresults(r, fn, args)

        if args.moduli or args.ssh_lines or args.crt_lines or args.dkim:
            if lcount == 0:
                _printresults({"type": "notfound", "results": {}}, fn, args)

        if fn != "-":
            f.close()

    sys.exit(_getret())
