pw_tls_client/py/pw_tls_client/generate_test_data.py - pigweed/pigweed - Git at Google

 # Copyright 2021 The Pigweed Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
 # use this file except in compliance with the License. You may obtain a copy of
 # the License at
 #
 #     https://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 # License for the specific language governing permissions and limitations under
 # the License.
 """Generate test data

 Generate data needed for unit tests, i.e. certificates, keys, and CRLSet.
 """

 import argparse
 import subprocess
 import sys
 from datetime import datetime, timedelta
 from typing import List, Tuple

 from cryptography import x509
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
 from cryptography.x509.oid import NameOID

 CERTS_AND_KEYS_HEADER = """// Copyright 2021 The Pigweed Authors
 //
 // Licensed under the Apache License, Version 2.0 (the "License"); you may not
 // use this file except in compliance with the License. You may obtain a copy
 // of the License at
 //
 //     https://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 // License for the specific language governing permissions and limitations under
 // the License.

 #pragma once

 #include "pw_bytes/span.h"

 """


 class Subject:
     """A subject wraps a name, private key and extensions for issuers
     to issue its certificate"""
     def __init__(self, name: str, extensions: List[Tuple[x509.ExtensionType,
                                                          bool]]):
         self._subject_name = x509.Name([
             x509.NameAttribute(NameOID.COUNTRY_NAME, u"US"),
             x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, u"California"),
             x509.NameAttribute(NameOID.LOCALITY_NAME, u"Mountain View"),
             x509.NameAttribute(NameOID.ORGANIZATION_NAME, name),
             x509.NameAttribute(NameOID.COMMON_NAME, u"Google-Pigweed"),
         ])
         self._private_key = rsa.generate_private_key(public_exponent=65537,
                                                      key_size=2048)
         self._extensions = extensions

     def subject_name(self) -> x509.Name:
         """Returns the subject name"""
         return self._subject_name

     def public_key(self) -> rsa.RSAPublicKey:
         """Returns the public key of this subject"""
         return self._private_key.public_key()

     def private_key(self) -> rsa.RSAPrivateKey:
         """Returns the private key of this subject"""
         return self._private_key

     def extensions(self) -> List[Tuple[x509.ExtensionType, bool]]:
         """Returns the requested extensions for issuer"""
         return self._extensions


 class CA(Subject):
     """A CA/Sub-ca that issues certificates"""
     def __init__(self, *args, **kwargs):
         ext = [(x509.BasicConstraints(True, None), True),
                (x509.KeyUsage(
                    digital_signature=False,
                    content_commitment=False,
                    key_encipherment=False,
                    data_encipherment=False,
                    key_agreement=False,
                    crl_sign=False,
                    encipher_only=False,
                    decipher_only=False,
                    key_cert_sign=True,
                ), True)]
         super().__init__(*args, extensions=ext, **kwargs)

     def sign(self, subject: Subject, not_before: datetime,
              not_after: datetime) -> x509.Certificate:
         """Issues a certificate for another CA/Sub-ca/Server"""
         builder = x509.CertificateBuilder()

         # Subject name is the target's subject name
         builder = builder.subject_name(subject.subject_name())

         # Issuer name is this CA/sub-ca's subject name
         builder = builder.issuer_name(self._subject_name)

         # Public key is the target's public key.
         builder = builder.public_key(subject.public_key())

         # Validity period.
         builder = builder.not_valid_before(not_before).not_valid_after(
             not_after)

         # Uses a random serial number
         builder = builder.serial_number(x509.random_serial_number())

         # Add extensions
         for extension, critical in subject.extensions():
             builder = builder.add_extension(extension, critical)

         # Sign and returns the certificate.
         return builder.sign(self._private_key, hashes.SHA256())

     def self_sign(self, not_before: datetime,
                   not_after: datetime) -> x509.Certificate:
         """Issues a self sign certificate"""
         return self.sign(self, not_before, not_after)


 class Server(Subject):
     """The end-entity server"""
     def __init__(self, *args, **kwargs):
         ext = [
             (x509.BasicConstraints(False, None), True),
             (x509.KeyUsage(
                 digital_signature=True,
                 content_commitment=False,
                 key_encipherment=False,
                 data_encipherment=False,
                 key_agreement=False,
                 crl_sign=False,
                 encipher_only=False,
                 decipher_only=False,
                 key_cert_sign=False,
             ), True),
             (x509.ExtendedKeyUsage([x509.ExtendedKeyUsageOID.SERVER_AUTH]),
              True),
         ]
         super().__init__(*args, extensions=ext, **kwargs)


 def c_escaped_string(data: bytes):
     """Generates a C byte string representation for a byte array

     For example, given a byte sequence of [0x12, 0x34, 0x56]. The function
     generates the following byte string code:

             {"\x12\x34\x56", 3}
     """
     body = ''.join([f'\\x{b:02x}' for b in data])
     return f'{{\"{body}\", {len(data)}}}'


 def byte_array_declaration(data: bytes, name: str) -> str:
     """Generates a ConstByteSpan declaration for a byte array"""
     type_name = '[[maybe_unused]] const pw::ConstByteSpan'
     array_body = f'pw::as_bytes(pw::span{c_escaped_string(data)})'
     return f'{type_name} {name} = {array_body};'


 class Codegen:
     """Base helper class for code generation"""
     def generate_code(self) -> str:
         """Generates C++ code for this object"""


 class PrivateKeyGen(Codegen):
     """Codegen class for a private key"""
     def __init__(self, key: rsa.RSAPrivateKey, name: str):
         self._key = key
         self._name = name

     def generate_code(self) -> str:
         """Code generation"""
         return byte_array_declaration(
             self._key.private_bytes(
                 serialization.Encoding.DER,
                 serialization.PrivateFormat.TraditionalOpenSSL,
                 serialization.NoEncryption()), self._name)


 class CertificateGen(Codegen):
     """Codegen class for a single certificate"""
     def __init__(self, cert: x509.Certificate, name: str):
         self._cert = cert
         self._name = name

     def generate_code(self) -> str:
         """Code generation"""
         return byte_array_declaration(
             self._cert.public_bytes(serialization.Encoding.DER), self._name)


 def generate_test_data() -> str:
     """Generates test data"""
     subjects: List[Codegen] = []

     # Working valid period.
     # Start from yesterday, to make sure we are in the valid period.
     not_before = datetime.utcnow() - timedelta(days=1)
     # Valid for 1 year.
     not_after = not_before + timedelta(days=365)

     # Generate a root-A CA certificates
     root_a = CA("root-A")
     subjects.append(
         CertificateGen(root_a.self_sign(not_before, not_after), "kRootACert"))

     # Generate a sub CA certificate signed by root-A.
     sub = CA("sub")
     subjects.append(
         CertificateGen(root_a.sign(sub, not_before, not_after), "kSubCACert"))

     # Generate a valid server certificate signed by sub
     server = Server("server")
     subjects.append(
         CertificateGen(sub.sign(server, not_before, not_after), "kServerCert"))
     subjects.append(PrivateKeyGen(server.private_key(), "kServerKey"))

     root_b = CA("root-B")
     subjects.append(
         CertificateGen(root_b.self_sign(not_before, not_after), "kRootBCert"))

     code = 'namespace {\n\n'
     for subject in subjects:
         code += subject.generate_code() + '\n\n'
     code += '}\n'

     return code


 def clang_format(file):
     subprocess.run([
         "clang-format",
         "-i",
         file,
     ], check=True)


 def parse_args():
     """Setup argparse."""
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "certs_and_keys_header",
         help="output header file for test certificates and keys")
     return parser.parse_args()


 def main() -> int:
     """Main"""
     args = parse_args()

     certs_and_keys = generate_test_data()

     with open(args.certs_and_keys_header, 'w') as header:
         header.write(CERTS_AND_KEYS_HEADER)
         header.write(certs_and_keys)

     clang_format(args.certs_and_keys_header)
     return 0


 if __name__ == "__main__":
     sys.exit(main())
	# Copyright 2021 The Pigweed Authors
	#
	# Licensed under the Apache License, Version 2.0 (the "License"); you may not
	# use this file except in compliance with the License. You may obtain a copy of
	# the License at
	#
	# https://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
	# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
	# License for the specific language governing permissions and limitations under
	# the License.
	"""Generate test data

	Generate data needed for unit tests, i.e. certificates, keys, and CRLSet.
	"""

	import argparse
	import subprocess
	import sys
	from datetime import datetime, timedelta
	from typing import List, Tuple

	from cryptography import x509
	from cryptography.hazmat.primitives import hashes
	from cryptography.hazmat.primitives import serialization
	from cryptography.hazmat.primitives.asymmetric import rsa
	from cryptography.x509.oid import NameOID

	CERTS_AND_KEYS_HEADER = """// Copyright 2021 The Pigweed Authors
	//
	// Licensed under the Apache License, Version 2.0 (the "License"); you may not
	// use this file except in compliance with the License. You may obtain a copy
	// of the License at
	//
	// https://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
	// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
	// License for the specific language governing permissions and limitations under
	// the License.

	#pragma once

	#include "pw_bytes/span.h"

	"""


	class Subject:
	"""A subject wraps a name, private key and extensions for issuers
	to issue its certificate"""
	def __init__(self, name: str, extensions: List[Tuple[x509.ExtensionType,
	bool]]):
	self._subject_name = x509.Name([
	x509.NameAttribute(NameOID.COUNTRY_NAME, u"US"),
	x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, u"California"),
	x509.NameAttribute(NameOID.LOCALITY_NAME, u"Mountain View"),
	x509.NameAttribute(NameOID.ORGANIZATION_NAME, name),
	x509.NameAttribute(NameOID.COMMON_NAME, u"Google-Pigweed"),
	])
	self._private_key = rsa.generate_private_key(public_exponent=65537,
	key_size=2048)
	self._extensions = extensions

	def subject_name(self) -> x509.Name:
	"""Returns the subject name"""
	return self._subject_name

	def public_key(self) -> rsa.RSAPublicKey:
	"""Returns the public key of this subject"""
	return self._private_key.public_key()

	def private_key(self) -> rsa.RSAPrivateKey:
	"""Returns the private key of this subject"""
	return self._private_key

	def extensions(self) -> List[Tuple[x509.ExtensionType, bool]]:
	"""Returns the requested extensions for issuer"""
	return self._extensions


	class CA(Subject):
	"""A CA/Sub-ca that issues certificates"""
	def __init__(self, args, *kwargs):
	ext = [(x509.BasicConstraints(True, None), True),
	(x509.KeyUsage(
	digital_signature=False,
	content_commitment=False,
	key_encipherment=False,
	data_encipherment=False,
	key_agreement=False,
	crl_sign=False,
	encipher_only=False,
	decipher_only=False,
	key_cert_sign=True,
	), True)]
	super().__init__(args, extensions=ext, *kwargs)

	def sign(self, subject: Subject, not_before: datetime,
	not_after: datetime) -> x509.Certificate:
	"""Issues a certificate for another CA/Sub-ca/Server"""
	builder = x509.CertificateBuilder()

	# Subject name is the target's subject name
	builder = builder.subject_name(subject.subject_name())

	# Issuer name is this CA/sub-ca's subject name
	builder = builder.issuer_name(self._subject_name)

	# Public key is the target's public key.
	builder = builder.public_key(subject.public_key())

	# Validity period.
	builder = builder.not_valid_before(not_before).not_valid_after(
	not_after)

	# Uses a random serial number
	builder = builder.serial_number(x509.random_serial_number())

	# Add extensions
	for extension, critical in subject.extensions():
	builder = builder.add_extension(extension, critical)

	# Sign and returns the certificate.
	return builder.sign(self._private_key, hashes.SHA256())

	def self_sign(self, not_before: datetime,
	not_after: datetime) -> x509.Certificate:
	"""Issues a self sign certificate"""
	return self.sign(self, not_before, not_after)


	class Server(Subject):
	"""The end-entity server"""
	def __init__(self, args, *kwargs):
	ext = [
	(x509.BasicConstraints(False, None), True),
	(x509.KeyUsage(
	digital_signature=True,
	content_commitment=False,
	key_encipherment=False,
	data_encipherment=False,
	key_agreement=False,
	crl_sign=False,
	encipher_only=False,
	decipher_only=False,
	key_cert_sign=False,
	), True),
	(x509.ExtendedKeyUsage([x509.ExtendedKeyUsageOID.SERVER_AUTH]),
	True),
	]
	super().__init__(args, extensions=ext, *kwargs)


	def c_escaped_string(data: bytes):
	"""Generates a C byte string representation for a byte array

	For example, given a byte sequence of [0x12, 0x34, 0x56]. The function
	generates the following byte string code:

	{"\x12\x34\x56", 3}
	"""
	body = ''.join([f'\\x{b:02x}' for b in data])
	return f'{{\"{body}\", {len(data)}}}'


	def byte_array_declaration(data: bytes, name: str) -> str:
	"""Generates a ConstByteSpan declaration for a byte array"""
	type_name = '[[maybe_unused]] const pw::ConstByteSpan'
	array_body = f'pw::as_bytes(pw::span{c_escaped_string(data)})'
	return f'{type_name} {name} = {array_body};'


	class Codegen:
	"""Base helper class for code generation"""
	def generate_code(self) -> str:
	"""Generates C++ code for this object"""


	class PrivateKeyGen(Codegen):
	"""Codegen class for a private key"""
	def __init__(self, key: rsa.RSAPrivateKey, name: str):
	self._key = key
	self._name = name

	def generate_code(self) -> str:
	"""Code generation"""
	return byte_array_declaration(
	self._key.private_bytes(
	serialization.Encoding.DER,
	serialization.PrivateFormat.TraditionalOpenSSL,
	serialization.NoEncryption()), self._name)


	class CertificateGen(Codegen):
	"""Codegen class for a single certificate"""
	def __init__(self, cert: x509.Certificate, name: str):
	self._cert = cert
	self._name = name

	def generate_code(self) -> str:
	"""Code generation"""
	return byte_array_declaration(
	self._cert.public_bytes(serialization.Encoding.DER), self._name)


	def generate_test_data() -> str:
	"""Generates test data"""
	subjects: List[Codegen] = []

	# Working valid period.
	# Start from yesterday, to make sure we are in the valid period.
	not_before = datetime.utcnow() - timedelta(days=1)
	# Valid for 1 year.
	not_after = not_before + timedelta(days=365)

	# Generate a root-A CA certificates
	root_a = CA("root-A")
	subjects.append(
	CertificateGen(root_a.self_sign(not_before, not_after), "kRootACert"))

	# Generate a sub CA certificate signed by root-A.
	sub = CA("sub")
	subjects.append(
	CertificateGen(root_a.sign(sub, not_before, not_after), "kSubCACert"))

	# Generate a valid server certificate signed by sub
	server = Server("server")
	subjects.append(
	CertificateGen(sub.sign(server, not_before, not_after), "kServerCert"))
	subjects.append(PrivateKeyGen(server.private_key(), "kServerKey"))

	root_b = CA("root-B")
	subjects.append(
	CertificateGen(root_b.self_sign(not_before, not_after), "kRootBCert"))

	code = 'namespace {\n\n'
	for subject in subjects:
	code += subject.generate_code() + '\n\n'
	code += '}\n'

	return code


	def clang_format(file):
	subprocess.run([
	"clang-format",
	"-i",
	file,
	], check=True)


	def parse_args():
	"""Setup argparse."""
	parser = argparse.ArgumentParser()
	parser.add_argument(
	"certs_and_keys_header",
	help="output header file for test certificates and keys")
	return parser.parse_args()


	def main() -> int:
	"""Main"""
	args = parse_args()

	certs_and_keys = generate_test_data()

	with open(args.certs_and_keys_header, 'w') as header:
	header.write(CERTS_AND_KEYS_HEADER)
	header.write(certs_and_keys)

	clang_format(args.certs_and_keys_header)
	return 0


	if __name__ == "__main__":
	sys.exit(main())