xref: /aosp_15_r20/external/pytorch/.ci/pytorch/create_test_cert.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1from datetime import datetime, timedelta
2from tempfile import mkdtemp
3
4from cryptography import x509
5from cryptography.hazmat.primitives import hashes, serialization
6from cryptography.hazmat.primitives.asymmetric import rsa
7from cryptography.x509.oid import NameOID
8
9
10temp_dir = mkdtemp()
11print(temp_dir)
12
13
14def genrsa(path):
15    key = rsa.generate_private_key(
16        public_exponent=65537,
17        key_size=2048,
18    )
19    with open(path, "wb") as f:
20        f.write(
21            key.private_bytes(
22                encoding=serialization.Encoding.PEM,
23                format=serialization.PrivateFormat.TraditionalOpenSSL,
24                encryption_algorithm=serialization.NoEncryption(),
25            )
26        )
27    return key
28
29
30def create_cert(path, C, ST, L, O, key):
31    subject = issuer = x509.Name(
32        [
33            x509.NameAttribute(NameOID.COUNTRY_NAME, C),
34            x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, ST),
35            x509.NameAttribute(NameOID.LOCALITY_NAME, L),
36            x509.NameAttribute(NameOID.ORGANIZATION_NAME, O),
37        ]
38    )
39    cert = (
40        x509.CertificateBuilder()
41        .subject_name(subject)
42        .issuer_name(issuer)
43        .public_key(key.public_key())
44        .serial_number(x509.random_serial_number())
45        .not_valid_before(datetime.utcnow())
46        .not_valid_after(
47            # Our certificate will be valid for 10 days
48            datetime.utcnow()
49            + timedelta(days=10)
50        )
51        .add_extension(
52            x509.BasicConstraints(ca=True, path_length=None),
53            critical=True,
54        )
55        .sign(key, hashes.SHA256())
56    )
57    # Write our certificate out to disk.
58    with open(path, "wb") as f:
59        f.write(cert.public_bytes(serialization.Encoding.PEM))
60    return cert
61
62
63def create_req(path, C, ST, L, O, key):
64    csr = (
65        x509.CertificateSigningRequestBuilder()
66        .subject_name(
67            x509.Name(
68                [
69                    # Provide various details about who we are.
70                    x509.NameAttribute(NameOID.COUNTRY_NAME, C),
71                    x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, ST),
72                    x509.NameAttribute(NameOID.LOCALITY_NAME, L),
73                    x509.NameAttribute(NameOID.ORGANIZATION_NAME, O),
74                ]
75            )
76        )
77        .sign(key, hashes.SHA256())
78    )
79    with open(path, "wb") as f:
80        f.write(csr.public_bytes(serialization.Encoding.PEM))
81    return csr
82
83
84def sign_certificate_request(path, csr_cert, ca_cert, private_ca_key):
85    cert = (
86        x509.CertificateBuilder()
87        .subject_name(csr_cert.subject)
88        .issuer_name(ca_cert.subject)
89        .public_key(csr_cert.public_key())
90        .serial_number(x509.random_serial_number())
91        .not_valid_before(datetime.utcnow())
92        .not_valid_after(
93            # Our certificate will be valid for 10 days
94            datetime.utcnow()
95            + timedelta(days=10)
96            # Sign our certificate with our private key
97        )
98        .sign(private_ca_key, hashes.SHA256())
99    )
100    with open(path, "wb") as f:
101        f.write(cert.public_bytes(serialization.Encoding.PEM))
102    return cert
103
104
105ca_key = genrsa(temp_dir + "/ca.key")
106ca_cert = create_cert(
107    temp_dir + "/ca.pem",
108    "US",
109    "New York",
110    "New York",
111    "Gloo Certificate Authority",
112    ca_key,
113)
114
115pkey = genrsa(temp_dir + "/pkey.key")
116csr = create_req(
117    temp_dir + "/csr.csr",
118    "US",
119    "California",
120    "San Francisco",
121    "Gloo Testing Company",
122    pkey,
123)
124
125cert = sign_certificate_request(temp_dir + "/cert.pem", csr, ca_cert, ca_key)
126