import gnu.getopt.Getopt;
import java.io.IOException;
import java.security.Security;
import java.util.Map;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.Objects;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.asn1.*;
import org.bouncycastle.asn1.pkcs.Attribute;

/**
 * P9AttrGen: A simple BouncyCastle PKCS#9 attribute generator.  Currently
 * it is only written to support a small subset of attributes that were needed
 * for JDK-8239950, but it can easily be expanded to support others.  This
 * application requires the inclusion of the gnu-getopt and BouncyCastle
 * Java libraries.
 */
public class P9AttrGen {

    public static final Map<String, ASN1ObjectIdentifier> ATTR_MAP = Map.of(
        "unstructuredName", new ASN1ObjectIdentifier("1.2.840.113549.1.9.2"),
        "signingTime", new ASN1ObjectIdentifier("1.2.840.113549.1.9.5"),
        "challengePassword", new ASN1ObjectIdentifier("1.2.840.113549.1.9.7"),
        "unstructuredAddress", new ASN1ObjectIdentifier("1.2.840.113549.1.9.8")
    );


    public static void usage() {
        System.err.println("Usage: P9Attr -t <TYPE> <-v VALSPEC> ...");
        System.err.println("OPTIONS:");
        System.err.println("-t <TYPE>      Create PKCS#9 Attribute of ID <TYPE>");
        System.err.println("-v <VALSPEC>   Add a value in the form <ID>:<VALUE>");
        System.err.println("               (may be used multiple times)");
        System.err.println("ATTRIBUTE TYPES:");
        System.err.println("    unstructuredName, signingTime, " +
                "challengePassword, unstructuredAddress");
        System.err.println("STRING TYPES: (simple string value)");
        System.err.println("    bmp, ia5, print, t61, univ, utf8");
        System.err.println("TIME TYPES:");
        System.err.println("    utctime: YYMMDDHHMMSSZ");
        System.err.println("    gentime: YYYYMMDDHHMMSS[.f]Z");
    }

    public static void main(String[] args) throws Exception {
        int c;
        String p9Type = null;
        List<ASN1Encodable> values = new ArrayList<>();

        Security.addProvider(new BouncyCastleProvider());

        Getopt g = new Getopt("P9AttrGen", args, "Ht:v:");
        while ((c = g.getopt()) != -1) {
            switch (c) {
                case 't':           // PKCS#9 Attribute type
                    p9Type = g.getOptarg();
                    break;
                case 'v':           // Value specifier
                    // Each one should follow the format before we add it
                    String[] idAndVal = g.getOptarg().split(":");
                    if (idAndVal.length != 2) {
                        System.err.println(
                                "Error: incorrect value specifier format");
                        System.exit(1);
                    } else {
                        values.add(getASN1Encodable(idAndVal));
                    }
                    break;
                case 'H':           // Help!
                    usage();
                    return;
                case '?':           // Unknown option or otherwise
                default:
                    usage();
                    System.exit(1);
            }
        }

        ASN1ObjectIdentifier p9oid = ATTR_MAP.get(Objects.requireNonNull(p9Type,
                    "A PKCS#9 type must be selected"));
        ASN1Set p9AttrSet = null;
        if (values.size() > 0) {
            p9AttrSet = new DERSet(values.toArray(new ASN1Encodable[0]));
        } else {
            throw new RuntimeException("Missing PKCS#9 attribute value(s)");
        }

        Attribute attr = new Attribute(p9oid, p9AttrSet);
        System.out.println(
                Base64.getMimeEncoder().encodeToString(attr.getEncoded()));
    }

    public static ASN1Encodable getASN1Encodable(String[] idAndVal)
            throws IOException {
        switch (idAndVal[0]) {
            case "bmp":             // BMPString
                return new DERBMPString(idAndVal[1]);
            case "gentime":         // GeneralizedTime
                return new ASN1GeneralizedTime(idAndVal[1]);
            case "ia5":             // IA5String
                return new DERIA5String(idAndVal[1]);
            case "print":           // PrintableString
                return new DERPrintableString(idAndVal[1]);
            case "t61":             // T61String/TeletexString
                return new DERT61String(idAndVal[1]);
            case "univ":            // UniversalString
                return new DERUniversalString(idAndVal[1].getBytes("UTF_32BE"));
            case "utctime":         // UTCTime
                return new ASN1UTCTime(idAndVal[1]);
            case "utf8":            // UTF8String
                return new DERUTF8String(idAndVal[1]);
            default:
                throw new RuntimeException("Unsupported id: " + idAndVal[0]);
        }
    }
}