package ui;

import model.asn1.ASN1Object;
import model.asn1.BitString;
import model.asn1.Bool;
import model.asn1.ObjectIdentifier;
import model.asn1.exceptions.ParseException;
import model.asn1.parsing.BytesReader;
import model.csr.CertificationRequest;
import model.pki.SubjectPublicKeyInfo;
import model.pki.cert.Certificate;
import model.pki.cert.Extension;
import model.pki.cert.TbsCertificate;

import java.util.Arrays;
import java.util.Base64;
import java.util.BitSet;

public class MgmtScreen implements UIHandler {
    private final JCA session;

    /**
     * EFFECTS: Init with the parent session.
     */
    public MgmtScreen(JCA session) {
        this.session = session;
    }

    @Override
    public void help() {
        System.out.print("show\tView the public key and CA certificate\n"
                + "csr\tGenerate a CSR for a upper-level CA to sign\n"
                + "install\tInstall a CA certificate\n"
                + "exit\tGo to main menu\n"
                + "help\tPrint this message\n");
    }

    /**
     * EFFECTS: Format the public key and CA
     */
    @Override
    public void show() {
        System.out.printf("Public Key:\t%s\n",
                Base64.getEncoder().encodeToString(session.getCa().getPublicKey().getEncoded()));
        if (!session.checkCA(true)) {
            return;
        }
        final TbsCertificate info = session.getCa().getCertificate().getCertificate();
        System.out.printf("Subject:\t%s\n", info.getSubject().toString());
        System.out.printf("Issuer:\t%s\n", info.getIssuer().toString());
        System.out.printf("Not Before:\t%s\n", info.getValidity().getNotBefore().getTimestamp());
        System.out.printf("Not After:\t%s\n", info.getValidity().getNotAfter().getTimestamp());
        System.out.printf("Signature:\t%s\n",
                Base64.getEncoder().encodeToString(Utils.byteToByte(info.getSubjectPublicKeyInfo()
                        .getSubjectPublicKey().getConvertedVal())));
    }

    private void handleCSR() {
        if (!session.checkCA(false)) {
            return;
        }
        try {
            CertificationRequest req = session.getCa().signCSR();
            System.out.println(Utils.toPEM(req.encodeDER(), "CERTIFICATE REQUEST"));
            session.log("Signed CA CSR.");
        } catch (Throwable e) {
            System.out.println(e.getMessage());
        }
    }

    private void validateCACertificateVersion(Certificate cert) throws ParseException {
        if (cert.getCertificate().getVersion() == null
                || cert.getCertificate().getVersion().getLong() != TbsCertificate.VERSION_V3) {
            throw new ParseException("The input certificate must be V3");
        }
    }

    private void validateCACertificatePublicKey(Certificate cert) throws ParseException {
        final SubjectPublicKeyInfo expectedPKInfo = session.getCa().getCAPublicKeyInfo();
        if (!Arrays.equals(cert.getCertificate().getSubjectPublicKeyInfo().getAlgorithm().getType().getInts(),
                expectedPKInfo.getAlgorithm().getType().getInts())
                || !Arrays.equals(cert.getCertificate().getSubjectPublicKeyInfo().getSubjectPublicKey().getVal(),
                expectedPKInfo.getSubjectPublicKey().getVal())) {
            throw new ParseException("The input certificate does not have the corresponding public key");
        }
    }

    private void validateCACertificateBasicConstraints(Certificate cert) throws ParseException {
        final Extension basicConstraints = cert.getCertificate().getExtension(ObjectIdentifier.OID_BASIC_CONSTRAINTS);
        if (basicConstraints == null
                || basicConstraints.getExtnValue().getBytes().length <= 0) {
            throw new ParseException("The certificate does not have a valid basicConstraints extension.");
        }
        final ASN1Object basicConstraintsValue =
                new ASN1Object(new BytesReader(basicConstraints.getExtnValue().getBytes()), false);
        if (basicConstraintsValue.getLength() <= 0) {
            throw new ParseException("The certificate does not have a valid basicConstraints extension.");
        }
        final ASN1Object bool =
                ASN1Object.parse(new BytesReader(basicConstraintsValue.encodeValueDER()), false);
        if (!(bool instanceof Bool)
                || !((Bool) bool).getValue()) {
            throw new ParseException("The certificate does not have a valid basicConstraints extension.");
        }
    }

    private void validateCACertificateKeyUsage(Certificate cert) throws ParseException {
        final Extension keyUsage = cert.getCertificate().getExtension(ObjectIdentifier.OID_KEY_USAGE);
        if (keyUsage == null
                || keyUsage.getExtnValue().getBytes().length <= 0) {
            throw new ParseException("The certificate does not have a valid keyUsage extension.");
        }
        final ASN1Object keyUsageValue =
                ASN1Object.parse(new BytesReader(keyUsage.getExtnValue().getBytes()), false);
        if (keyUsageValue.getLength() <= 0
                || !(keyUsageValue instanceof BitString)) {
            throw new ParseException("The certificate does not have a valid keyUsage extension.");
        }
        final BitSet bitSet = BitSet.valueOf(Utils.byteToByte(((BitString) keyUsageValue).getVal()));
        if (!bitSet.get(7) || !bitSet.get(2) || !bitSet.get(1)) {
            throw new ParseException("The certificate does not have a valid keyUsage extension.");
        }
    }

    private void handleInstall() {
        if (!session.checkCA(false)) {
            return;
        }
        try {
            final Byte[] in = session.handleInputPEM("CERTIFICATE");
            Certificate cert = new Certificate(new BytesReader(in), false);
            validateCACertificateVersion(cert);
            validateCACertificatePublicKey(cert);
            validateCACertificateBasicConstraints(cert);
            validateCACertificateKeyUsage(cert);
            session.getCa().installCertificate(cert);
            session.log("A CA certificate is installed.");
        } catch (ParseException e) {
            System.out.println(e.getMessage());
        }
    }

    @Override
    public void command(String... args) {
        switch (args[0]) {
            case "csr":
                handleCSR();
                break;
            case "install":
                handleInstall();
                break;
            default:
                help();
                break;
        }
    }

    /**
     * EFFECTS: Go to main menu
     */
    @Override
    public Screen exit() {
        return Screen.MAIN;
    }

    @Override
    public String getPS1() {
        return "/ca/ #";
    }
}