Make the CertificateBuilder interface more like the CSRBuilder

1 files changed, 92 insertions, 27 deletions

diff --git a/src/cryptography/x509.py b/src/cryptography/x509.py
index 7cb42f57..c04b8c9c 100644
--- a/src/cryptography/x509.py
+++ b/src/cryptography/x509.py
@@ -18,6 +18,7 @@ from six.moves import urllib_parse
 
 from cryptography import utils
 from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.asymmetric import dsa, ec, rsa
 
 
 _OID_NAMES = {
@@ -1598,89 +1599,153 @@ class CertificateSigningRequestBuilder(object):
 
 
 class CertificateBuilder(object):
-    def __init__(self):
+    def __init__(self, version=None, issuer_name=None, subject_name=None,
+                 public_key=None, serial_number=None, not_valid_before=None,
+                 not_valid_after=None, extensions=[]):
         """
         Creates an empty X.509 certificate (version 1).
         """
-        self._version = Version.v1
-        self._issuer_name = None
-        self._subject_name = None
-        self._public_key = None
-        self._serial_number = None
-        self._not_valid_before = None
-        self._not_valid_after = None
-        self._extensions = []
+        self._version = version
+        self._issuer_name = issuer_name
+        self._subject_name = subject_name
+        self._public_key = public_key
+        self._serial_number = serial_number
+        self._not_valid_before = not_valid_before
+        self._not_valid_after = not_valid_after
+        self._extensions = extensions
 
-    def set_version(self, version):
+    def version(self, version):
         """
         Sets the X.509 version required by decoders.
         """
         if not isinstance(version, Version):
             raise TypeError('Expecting x509.Version object.')
-        self._version = version
+        if self._version is not None:
+            raise ValueError('The version may only be set once.')
+        return CertificateBuilder(
+            version, self._issuer_name, self._subject_name, self._public_key,
+            self._serial_number, self._not_valid_before,
+            self._not_valid_after, self._extensions
+        )
 
-    def set_issuer_name(self, name):
+    def issuer_name(self, name):
         """
         Sets the CA's distinguished name.
         """
         if not isinstance(name, Name):
             raise TypeError('Expecting x509.Name object.')
-        self._issuer_name = name
+        if self._issuer_name is not None:
+            raise ValueError('The issuer name may only be set once.')
+        return CertificateBuilder(
+            self._version, name, self._subject_name, self._public_key,
+            self._serial_number, self._not_valid_before,
+            self._not_valid_after, self._extensions
+        )
 
-    def set_subject_name(self, name):
+    def subject_name(self, name):
         """
         Sets the requestor's distinguished name.
         """
         if not isinstance(name, Name):
             raise TypeError('Expecting x509.Name object.')
-        self._subject_name = name
+        if self._issuer_name is not None:
+            raise ValueError('The subject name may only be set once.')
+        return CertificateBuilder(
+            self._version, self._issuer_name, name, self._public_key,
+            self._serial_number, self._not_valid_before,
+            self._not_valid_after, self._extensions
+        )
 
-    def set_public_key(self, public_key):
+    def public_key(self, key):
         """
         Sets the requestor's public key (as found in the signing request).
         """
-        # TODO: check type.
-        self._public_key = public_key
+        if not isinstance(key, (dsa.DSAPublicKey, rsa.RSAPublicKey,
+                                ec.EllipticCurvePublicKey)):
+            raise TypeError('Expecting one of DSAPublicKey, RSAPublicKey,'
+                            ' or EllipticCurvePublicKey.')
+        if self._public_key is not None:
+            raise ValueError('The public key may only be set once.')
+        return CertificateBuilder(
+            self._version, self._issuer_name, self._subject_name, key,
+            self._serial_number, self._not_valid_before,
+            self._not_valid_after, self._extensions
+        )
 
-    def set_serial_number(self, serial_number):
+    def serial_number(self, number):
         """
         Sets the certificate serial number.
         """
-        if not isinstance(serial_number, six.integer_types):
+        if not isinstance(number, six.integer_types):
             raise TypeError('Serial number must be of integral type.')
-        self._serial_number = serial_number
+        if self._public_key is not None:
+            raise ValueError('The serial number may only be set once.')
+        return CertificateBuilder(
+            self._version, self._issuer_name, self._subject_name,
+            self._public_key, number, self._not_valid_before,
+            self._not_valid_after, self._extensions
+        )
 
-    def set_not_valid_before(self, time):
+    def not_valid_before(self, time):
         """
         Sets the certificate activation time.
         """
         # TODO: require UTC datetime?
         if not isinstance(time, datetime.datetime):
             raise TypeError('Expecting datetime object.')
-        self._not_valid_before = time
+        if self._not_valid_before is not None:
+            raise ValueError('The not valid before may only be set once.')
+        return CertificateBuilder(
+            self._version, self._issuer_name, self._subject_name,
+            self._public_key, self._serial_number, time,
+            self._not_valid_after, self._extensions
+        )
 
-    def set_not_valid_after(self, time):
+    def not_valid_after(self, time):
         """
         Sets the certificate expiration time.
         """
         # TODO: require UTC datetime?
         if not isinstance(time, datetime.datetime):
             raise TypeError('Expecting datetime object.')
-        self._not_valid_after = time
+        if self._not_valid_before is not None:
+            raise ValueError('The not valid after may only be set once.')
+        return CertificateBuilder(
+            self._version, self._issuer_name, self._subject_name,
+            self._public_key, self._serial_number, self._not_valid_before,
+            time, self._extensions
+        )
 
-    def add_extension(self, extension):
+    def add_extension(self, extension, critical):
         """
         Adds an X.509 extension to the certificate.
         """
+        if isinstance(extension, BasicConstraints):
+            extension = Extension(OID_BASIC_CONSTRAINTS, critical, extension)
+        elif isinstance(extension, SubjectAlternativeName):
+            extension = Extension(
+                OID_SUBJECT_ALTERNATIVE_NAME, critical, extension
+            )
+        else:
+            raise NotImplementedError('Unsupported X.509 extension.')
         if not isinstance(extension, Extension):
             raise TypeError('Expecting x509.Extension object.')
+
+        # TODO: This is quadratic in the number of extensions
         for e in self._extensions:
             if e.oid == extension.oid:
                 raise ValueError('This extension has already been set.')
-        self._extensions.append(extension)
+
+        return CertificateBuilder(
+            self._version, self._issuer_name, self._subject_name,
+            self._public_key, self._serial_number, self._not_valid_before,
+            self._not_valid_after, self._extensions + [extension]
+        )
 
     def sign(self, backend, private_key, algorithm):
         """
         Signs the certificate using the CA's private key.
         """
+        if self._version is None:
+            self._version = Version.v1
         return backend.sign_x509_certificate(self, private_key, algorithm)