package com.minamp.net;

import java.io.IOException;
import java.security.MessageDigest;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Locale;

import com.minamp.bc.asn1.ASN1Encoding;
import com.minamp.bc.asn1.ASN1ObjectIdentifier;
import com.minamp.bc.asn1.ASN1OctetString;
import com.minamp.bc.asn1.pkcs.PKCSObjectIdentifiers;
import com.minamp.bc.asn1.x500.RDN;
import com.minamp.bc.asn1.x500.X500Name;
import com.minamp.bc.asn1.x500.style.BCStyle;
import com.minamp.bc.asn1.x500.style.IETFUtils;
import com.minamp.bc.asn1.x509.BasicConstraints;
import com.minamp.bc.asn1.x509.Certificate;
import com.minamp.bc.asn1.x509.Extension;
import com.minamp.bc.asn1.x509.Extensions;
import com.minamp.bc.asn1.x509.GeneralName;
import com.minamp.bc.asn1.x509.GeneralNames;
import com.minamp.bc.asn1.x9.X9ObjectIdentifiers;
import com.minamp.bc.crypto.Digest;
import com.minamp.bc.crypto.digests.SHA1Digest;
import com.minamp.bc.crypto.digests.SHA256Digest;
import com.minamp.bc.crypto.digests.SHA384Digest;
import com.minamp.bc.crypto.digests.SHA512Digest;
import com.minamp.bc.crypto.params.AsymmetricKeyParameter;
import com.minamp.bc.crypto.signers.DSADigestSigner;
import com.minamp.bc.crypto.signers.ECDSASigner;
import com.minamp.bc.crypto.signers.RSADigestSigner;
import com.minamp.bc.crypto.util.PublicKeyFactory;

/**
 * Minimal Web PKI path validation for the bundled TLS stack:
 * hostname match, validity window, signature chain up to a bundled Mozilla root.
 * Optionally one pinned leaf certificate (SHA-256) is trusted as-is, for self-signed servers.
 * Shared with Minamp (tools/tlskit in the iOS ports' history); keep the copies in sync.
 */
public final class CertVerifier {

    /** Texts shown to the user; apps replace them with their own language. Formats take the host. */
    public static String msgMismatch = "The server certificate is for a different address (%1$s).";
    public static String msgDates = "The server certificate has expired, or the phone's date is wrong.";

    /** Thrown when the certificate is neither CA-trusted nor the pinned one: the app can offer to trust it. */
    public static final class Untrusted extends IOException {
        public final String fingerprint, subject, issuer, expires;
        public final boolean changed; // a pin was set and this is a different certificate
        public final List<byte[]> chain; // as the server sent it, leaf first

        Untrusted(String fp, String subject, String issuer, String expires, boolean changed, List<byte[]> chain) {
            super("untrusted certificate " + fp);
            this.chain = chain;
            this.fingerprint = fp;
            this.subject = subject;
            this.issuer = issuer;
            this.expires = expires;
            this.changed = changed;
        }
    }

    private static volatile String pin = "";

    /** SHA-256 of the one leaf certificate to accept without a CA (hex, separators ignored); "" for none. */
    public static void setPin(String p) {
        String n = norm(p);
        if (!n.equals(pin)) {
            pin = n;
            Http.forgetSessions();
        }
    }

    public static String norm(String s) {
        if (s == null) return "";
        return s.replaceAll("[^0-9A-Fa-f]", "").toUpperCase(Locale.US);
    }

    /** "AB:CD:..." SHA-256 of a DER certificate. */
    public static String sha256(byte[] der) {
        try {
            byte[] d = MessageDigest.getInstance("SHA-256").digest(der);
            StringBuilder b = new StringBuilder();
            for (int i = 0; i < d.length; i++) {
                if (i > 0) b.append(':');
                String x = Integer.toHexString(d[i] & 0xff).toUpperCase(Locale.US);
                if (x.length() == 1) b.append('0');
                b.append(x);
            }
            return b.toString();
        } catch (Exception e) {
            return "";
        }
    }

    private static Untrusted untrusted(List<byte[]> chainDer, boolean changed) {
        byte[] leafDer = chainDer.get(0);
        String subject = "?", issuer = "?", expires = "?";
        try {
            Certificate c = Certificate.getInstance(leafDer);
            subject = String.valueOf(c.getSubject());
            issuer = String.valueOf(c.getIssuer());
            expires = new java.text.SimpleDateFormat("yyyy-MM-dd", Locale.US).format(c.getEndDate().getDate());
        } catch (Exception ignored) {}
        return new Untrusted(sha256(leafDer), subject, issuer, expires, changed, chainDer);
    }

    /** Lookup of trusted root certificates (DER) by the DER encoding of their subject name. */
    public interface RootStore {
        List<byte[]> find(byte[] subjectDer);
    }

    private static RootStore roots;
    /** Chains already proven to lead to a trusted root (keyed by a hash of the chain), so a
     *  reconnect only re-checks the host name and dates, not every signature again. */
    private static final java.util.Set<String> verifiedChains = new java.util.HashSet<String>();

    /** Changing what's trusted invalidates everything trusted so far: cached chains and TLS sessions. */
    public static void setRootStore(RootStore store) {
        roots = store;
        synchronized (verifiedChains) { verifiedChains.clear(); }
        Http.forgetSessions();
    }

    private CertVerifier() {}


    public static void verify(List<byte[]> chainDer, String host) throws IOException {
        if (chainDer == null || chainDer.isEmpty()) throw new IOException("Server sent no certificate");
        String p = pin;
        if (p.length() > 0) {
            // a pinned server: exactly that certificate, whatever its name or issuer
            if (norm(sha256(chainDer.get(0))).equals(p)) return;
            throw untrusted(chainDer, true);
        }
        if (roots == null) throw new IOException("No trust store configured");
        java.io.ByteArrayOutputStream all = new java.io.ByteArrayOutputStream();
        for (int i = 0; i < chainDer.size(); i++) all.write(chainDer.get(i), 0, chainDer.get(i).length);
        String chainKey = sha1Hex(all.toByteArray());
        boolean known;
        synchronized (verifiedChains) { known = verifiedChains.contains(chainKey); }
        if (known) {
            Certificate leaf = Certificate.getInstance(chainDer.get(0));
            if (!hostMatches(leaf, host)) throw new IOException(String.format(msgMismatch, host));
            if (new Date().getTime() > 1704067200000L) checkDates(leaf, new Date());
            return;
        }
        verifyFully(chainDer, host);
        synchronized (verifiedChains) { verifiedChains.add(chainKey); }
    }

    private static void verifyFully(List<byte[]> chainDer, String host) throws IOException {
        final List<byte[]> chain0 = chainDer;
        List<Certificate> chain = new ArrayList<Certificate>();
        for (int i = 0; i < chainDer.size(); i++) chain.add(Certificate.getInstance(chainDer.get(i)));

        Certificate leaf = chain.get(0);
        if (!hostMatches(leaf, host)) {
            throw new IOException(String.format(msgMismatch, host));
        }

        Date now = new Date();
        // An unset clock on an old phone (e.g. 1980 after a battery pull) would make every
        // certificate look "not yet valid"; only enforce dates when the clock looks sane.
        boolean clockSane = now.getTime() > 1704067200000L; // 2024-01-01

        Certificate current = leaf;
        for (int depth = 0; depth < 8; depth++) {
            if (clockSane) checkDates(current, now);

            byte[] issuerDer = current.getIssuer().getEncoded(ASN1Encoding.DER);

            // 1. Anchored directly by a trusted root?
            List<byte[]> anchors = roots.find(issuerDer);
            if (anchors != null) {
                byte[] currentDer = current.getEncoded(ASN1Encoding.DER);
                for (int i = 0; i < anchors.size(); i++) {
                    // the very certificate the user chose to trust (e.g. self-signed)
                    if (java.util.Arrays.equals(anchors.get(i), currentDer)) return;
                    Certificate root = Certificate.getInstance(anchors.get(i));
                    // anything that signs must be a CA: a trusted site's own certificate can't vouch for another
                    if (isCa(root) && signatureValid(current, root)) return; // trusted
                }
            }

            // 2. Otherwise find the issuing intermediate among what the server sent.
            Certificate next = null;
            for (int i = 1; i < chain.size(); i++) {
                Certificate cand = chain.get(i);
                if (cand == current) continue;
                if (!java.util.Arrays.equals(cand.getSubject().getEncoded(ASN1Encoding.DER), issuerDer)) continue;
                if (!isCa(cand)) continue;
                if (signatureValid(current, cand)) { next = cand; break; }
            }
            if (next == null) throw untrusted(chain0, false);
            current = next;
        }
        throw new IOException("Certificate chain too long");
    }

    private static void checkDates(Certificate c, Date now) throws IOException {
        Date from = c.getStartDate().getDate();
        Date to = c.getEndDate().getDate();
        if (now.before(from) || now.after(to)) {
            throw new IOException(msgDates);
        }
    }

    private static boolean isCa(Certificate c) {
        Extensions ext = c.getTBSCertificate().getExtensions();
        if (ext == null) return c.getTBSCertificate().getVersionNumber() < 3; // v1 roots/intermediates
        BasicConstraints bc = BasicConstraints.fromExtensions(ext);
        return bc != null && bc.isCA();
    }

    static boolean signatureValid(Certificate subject, Certificate issuer) {
        try {
            ASN1ObjectIdentifier alg = subject.getSignatureAlgorithm().getAlgorithm();
            byte[] tbs = subject.getTBSCertificate().getEncoded(ASN1Encoding.DER);
            byte[] sig = subject.getSignature().getOctets();
            AsymmetricKeyParameter key = PublicKeyFactory.createKey(issuer.getSubjectPublicKeyInfo());

            if (alg.equals(PKCSObjectIdentifiers.sha256WithRSAEncryption)) return rsa(new SHA256Digest(), key, tbs, sig);
            if (alg.equals(PKCSObjectIdentifiers.sha384WithRSAEncryption)) return rsa(new SHA384Digest(), key, tbs, sig);
            if (alg.equals(PKCSObjectIdentifiers.sha512WithRSAEncryption)) return rsa(new SHA512Digest(), key, tbs, sig);
            if (alg.equals(PKCSObjectIdentifiers.sha1WithRSAEncryption)) return rsa(new SHA1Digest(), key, tbs, sig);
            if (alg.equals(X9ObjectIdentifiers.ecdsa_with_SHA256)) return ecdsa(new SHA256Digest(), key, tbs, sig);
            if (alg.equals(X9ObjectIdentifiers.ecdsa_with_SHA384)) return ecdsa(new SHA384Digest(), key, tbs, sig);
            if (alg.equals(X9ObjectIdentifiers.ecdsa_with_SHA512)) return ecdsa(new SHA512Digest(), key, tbs, sig);
            return false;
        } catch (Exception e) {
            return false;
        }
    }

    private static boolean rsa(Digest d, AsymmetricKeyParameter key, byte[] tbs, byte[] sig) {
        RSADigestSigner s = new RSADigestSigner(d);
        s.init(false, key);
        s.update(tbs, 0, tbs.length);
        return s.verifySignature(sig);
    }

    private static boolean ecdsa(Digest d, AsymmetricKeyParameter key, byte[] tbs, byte[] sig) {
        DSADigestSigner s = new DSADigestSigner(new ECDSASigner(), d);
        s.init(false, key);
        s.update(tbs, 0, tbs.length);
        return s.verifySignature(sig);
    }

    // ---- hostname checking ----

    static boolean hostMatches(Certificate leaf, String host) {
        host = host.toLowerCase(Locale.US);
        boolean ip = isIpLiteral(host);
        boolean sawDns = false;
        Extensions ext = leaf.getTBSCertificate().getExtensions();
        if (ext != null) {
            GeneralNames names = GeneralNames.fromExtensions(ext, Extension.subjectAlternativeName);
            if (names != null) {
                GeneralName[] arr = names.getNames();
                for (int i = 0; i < arr.length; i++) {
                    GeneralName n = arr[i];
                    if (n.getTagNo() == GeneralName.dNSName) {
                        sawDns = true;
                        // String.valueOf, not getName().toString(): Dalvik < 4.0 can't call Object methods via an interface
                        if (!ip && dnsMatch(String.valueOf(n.getName()), host)) return true;
                    } else if (n.getTagNo() == GeneralName.iPAddress && ip) {
                        byte[] a = ASN1OctetString.getInstance(n.getName()).getOctets();
                        if (ipString(a).equals(host)) return true;
                    }
                }
            }
        }
        if (!sawDns && !ip) {
            X500Name subject = X500Name.getInstance(leaf.getSubject());
            RDN[] cns = subject.getRDNs(BCStyle.CN);
            for (int i = 0; i < cns.length; i++) {
                String cn = IETFUtils.valueToString(cns[i].getFirst().getValue());
                if (dnsMatch(cn, host)) return true;
            }
        }
        return false;
    }

    static boolean dnsMatch(String pattern, String host) {
        pattern = pattern.toLowerCase(Locale.US);
        if (pattern.endsWith(".")) pattern = pattern.substring(0, pattern.length() - 1);
        if (host.endsWith(".")) host = host.substring(0, host.length() - 1);
        if (pattern.startsWith("*.")) {
            String suffix = pattern.substring(1); // ".example.com"
            if (!host.endsWith(suffix)) return false;
            String label = host.substring(0, host.length() - suffix.length());
            return label.length() > 0 && label.indexOf('.') < 0;
        }
        return pattern.equals(host);
    }

    static boolean isIpLiteral(String host) {
        if (host.indexOf(':') >= 0) return true;
        for (int i = 0; i < host.length(); i++) {
            char c = host.charAt(i);
            if (!(c == '.' || (c >= '0' && c <= '9'))) return false;
        }
        return true;
    }

    private static String ipString(byte[] a) {
        if (a.length == 4) {
            return (a[0] & 255) + "." + (a[1] & 255) + "." + (a[2] & 255) + "." + (a[3] & 255);
        }
        return "";
    }

    public static String sha1Hex(byte[] data) {
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-1");
            byte[] h = md.digest(data);
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < h.length; i++) {
                String s = Integer.toHexString(h[i] & 255);
                if (s.length() == 1) sb.append('0');
                sb.append(s);
            }
            return sb.toString();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}