1#!/usr/bin/env python
2#
3# Copyright (C) 2016 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the 'License');
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an 'AS IS' BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16import random
17
18from OpenSSL import crypto
19from Crypto.PublicKey import RSA
20
21class Certificate(object):
22    cert = None
23    key = None
24    def __init__(self, cert, key):
25        self.cert = cert
26        self.key = key
27
28    def cert_pem(self):
29        return crypto.dump_certificate(crypto.FILETYPE_PEM, self.cert)
30
31    def key_pem(self):
32        return crypto.dump_privatekey(crypto.FILETYPE_PEM, self.key)
33
34    def save_to_file(self, path):
35        with open(path, "w") as f:
36            f.write(self.cert_pem())
37            f.write(self.key_pem())
38
39    def save_cert_to_file(self, path):
40        with open(path, "w") as f:
41            f.write(self.cert_pem())
42
43    @staticmethod
44    def from_file(path):
45        with open(path) as f:
46            data = f.read()
47            cert = crypto.load_certificate(crypto.FILETYPE_PEM, data)
48            key = crypto.load_privatekey(crypto.FILETYPE_PEM, data)
49            return Certificate(cert, key)
50
51    @staticmethod
52    def create(cn, issuer=None, key=None, keysize=2048, digest="sha256",
53            notBefore="20150101000000+0000", notAfter="20300101000000+0000",
54            additional_extensions=None):
55        if key is None:
56            key = crypto.PKey()
57            key.generate_key(crypto.TYPE_RSA, keysize)
58
59        cert = crypto.X509()
60        cert.set_pubkey(key)
61        cert.set_version(2)
62        cert.set_serial_number(random.randint(0, 2**20))
63
64        cert.set_notBefore(notBefore)
65        cert.set_notAfter(notAfter)
66        cert.get_subject().CN = cn
67        cert.set_issuer(cert.get_subject() if issuer is None else issuer.cert.get_subject())
68        # Add the CA=True basic constraint
69        basicContraints = crypto.X509Extension("basicConstraints", True, "CA:TRUE")
70        cert.add_extensions([basicContraints])
71        if additional_extensions is not None:
72            cert.add_extensions(additional_extensions)
73
74        signing_key = key if issuer is None else issuer.key
75        cert.sign(signing_key, digest)
76
77        return Certificate(cert, key)
78
79if __name__ == "__main__":
80    # Generate test certificates.
81    a = Certificate.create("Root A")
82    a_sha1 = Certificate.create("Root A", key=a.key, digest="sha1")
83    b = Certificate.create("Root B")
84    a_to_b = Certificate.create("Root A", b, a.key)
85    b_to_a = Certificate.create("Root B", a, b.key)
86    leaf1 = Certificate.create("Leaf", a)
87    intermediate_a = Certificate.create("intermediate", a)
88    intermediate_b = Certificate.create("intermediate", b, intermediate_a.key)
89    leaf2 = Certificate.create("Leaf 2", intermediate_a)
90
91    # Save test certificates.
92    a.save_cert_to_file("a.pem")
93    a_sha1.save_cert_to_file("a_sha1.pem")
94    b.save_cert_to_file("b.pem")
95    a_to_b.save_cert_to_file("a_to_b.pem")
96    b_to_a.save_cert_to_file("b_to_a.pem")
97    leaf1.save_cert_to_file("leaf1.pem")
98    leaf2.save_cert_to_file("leaf2.pem")
99    intermediate_a.save_cert_to_file("intermediate_a.pem")
100    intermediate_b.save_cert_to_file("intermediate_b.pem")
101