From 52d85e47faa056b4e8d0c5c2007bf8887c00ffcc 2020-06-01 16:23:21
From: Branko Majic <branko@majic.rs>
Date: 2020-06-01 16:23:21
Subject: [PATCH] GC-37: Added support for requesting custom RSA key size when initialising the CA hierarchy:

- Added functional test.
- Added unit tests.
- Added new CLI option for specifying the algorithm.
- Implemented KeyGenerator factory-like class that can be called to
  generate a private key with desired specification.
- The init init function now accepts a callable that is used to
  generate private keys.
- The generate_ca_hierarchy function now accepts a callable that is
  used to generate private keys.
- Updated existing unit tests to cope with changes to the init
  and generate_ca_hierarchy function signatures.
- Updated existing unit tests to cope with changes to existing
  functionality.
- Updated existing functional tests to cope with changes in command
  output.

---

diff --git a/functional_tests/test_init.py b/functional_tests/test_init.py
index d6765d43d4fb02cce36aafcf15111560a9457252..a9d682562cba06a0c7b052408fd2bc630ea8516b 100644
--- a/functional_tests/test_init.py
+++ b/functional_tests/test_init.py
@@ -58,7 +58,7 @@ def test_initialisation_on_fresh_directory(tmpdir):
     # text to John that the directory has been initialised.
     assert exit_code == 0
     assert stderr == ""
-    assert "CA hierarchy initialised" in stdout
+    assert "CA hierarchy initialised using 2048-bit RSA keys." in stdout
 
     # The tool also points John to generated key and certificate material.
     assert ".gimmecert/ca/level1.key.pem" in stdout
@@ -149,7 +149,7 @@ def test_initialisation_with_custom_base_name(tmpdir):
     # his CA hierarchy has been initialised..
     assert exit_code == 0
     assert stderr == ""
-    assert "CA hierarchy initialised." in stdout
+    assert "CA hierarchy initialised using 2048-bit RSA keys." in stdout
 
     # Just before he starts using the CA certificates further, he
     # decides to double-check the results. He runs a couple of
@@ -193,7 +193,7 @@ def test_initialisation_with_custom_hierarchy_depth(tmpdir):
     # more CA artifacts listed now.
     assert exit_code == 0
     assert stderr == ""
-    assert "CA hierarchy initialised." in stdout
+    assert "CA hierarchy initialised using 2048-bit RSA keys." in stdout
     assert ".gimmecert/ca/level1.key.pem" in stdout
     assert ".gimmecert/ca/level1.cert.pem" in stdout
     assert ".gimmecert/ca/level2.key.pem" in stdout
@@ -284,3 +284,55 @@ def test_initialisation_with_custom_hierarchy_depth(tmpdir):
 
     # He is happy to see that verification succeeds.
     assert error_code == 0
+
+
+def test_initialisation_with_rsa_private_key_specificiation(tmpdir):
+    # John is looking into improving the security of one of his
+    # projects. Amongst other things, John is interested in using
+    # stronger private keys for his TLS services - which he wants to
+    # try out in his test envioronment first.
+
+    # John knows that the Gimmecert tool uses 2048-bit RSA keys for
+    # the CA hierarchy, but what he would really like to do is specify
+    # himself what kind of private key should be generated
+    # instead. He checks-out the help for the init command first.
+    stdout, _, _ = run_command('gimmecert', 'init', '-h')
+
+    # John noticies there is an option to provide a custom key
+    # specification to the tool, that he can specify the length of
+    # the RSA private keys, and that the default is "rsa:2048".
+    assert "--key-specification" in stdout
+    assert " -k" in stdout
+    assert "rsa:BIT_LENGTH" in stdout
+    assert "Default is rsa:2048" in stdout
+
+    # John switches to his project directory.
+    tmpdir.chdir()
+
+    # He initalises the CA hierarchy, requesting to use 4096-bit RSA
+    # keys.
+    stdout, stderr, exit_code = run_command('gimmecert', 'init', '--key-specification', 'rsa:4096')
+
+    # Command finishes execution with success, and John notices that
+    # the tool has informed him of what the private key algorithm is
+    # in use for the CA hierarchy.
+    assert exit_code == 0
+    assert stderr == ""
+    assert "CA hierarchy initialised using 4096-bit RSA keys." in stdout
+
+    # John goes ahead and inspects the CA private key to ensure his
+    # private key specification has been accepted.
+    stdout, stderr, exit_code = run_command('openssl', 'rsa', '-noout', '-text', '-in', '.gimmecert/ca/level1.key.pem')
+
+    assert exit_code == 0
+    assert stderr == ""
+    assert "Private-Key: (4096 bit)" in stdout
+
+    # John also does a quick check on the generated certificate's
+    # signing and public key algorithm.
+    stdout, stderr, exit_code = run_command('openssl', 'x509', '-noout', '-text', '-in', '.gimmecert/ca/level1.cert.pem')
+
+    assert exit_code == 0
+    assert stderr == ""
+    assert "Signature Algorithm: sha256WithRSAEncryption" in stdout
+    assert "Public-Key: (4096 bit)" in stdout
diff --git a/gimmecert/cli.py b/gimmecert/cli.py
index ea0612a2f09d1c2d7454291982050caa9573996b..413edf5f08fbbfd2ea90c3487d78529ff76cb1c0 100644
--- a/gimmecert/cli.py
+++ b/gimmecert/cli.py
@@ -25,8 +25,10 @@ import sys
 
 from .decorators import subcommand_parser, get_subcommand_parser_setup_functions
 from .commands import client, help_, init, renew, server, status, usage, ExitCode
+from .crypto import KeyGenerator
 
 
+ERROR_ARGUMENTS = 2
 ERROR_GENERIC = 10
 
 
@@ -81,13 +83,16 @@ def setup_init_subcommand_parser(parser, subparsers):
     subparser = subparsers.add_parser('init', description='Initialise CA hierarchy.')
     subparser.add_argument('--ca-base-name', '-b', help="Base name to use for CA naming. Default is to use the working directory base name.")
     subparser.add_argument('--ca-hierarchy-depth', '-d', type=int, help="Depth of CA hierarchy to generate. Default is 1", default=1)
+    subparser.add_argument('--key-specification', '-k', type=KeyGenerator,
+                           help='''Default specification/parameters to use for private key generation. \
+    For RSA keys, use format rsa:BIT_LENGTH. Default is rsa:2048.''', default="rsa:2048")
 
     def init_wrapper(args):
         project_directory = os.getcwd()
         if args.ca_base_name is None:
             args.ca_base_name = os.path.basename(project_directory)
 
-        return init(sys.stdout, sys.stderr, project_directory, args.ca_base_name, args.ca_hierarchy_depth)
+        return init(sys.stdout, sys.stderr, project_directory, args.ca_base_name, args.ca_hierarchy_depth, args.key_specification)
 
     subparser.set_defaults(func=init_wrapper)
 
diff --git a/gimmecert/commands.py b/gimmecert/commands.py
index 24f6fb32252d30c851b6b97c633939c9520c6012..319fed20c5f9cd7505690df37239945416d8813d 100644
--- a/gimmecert/commands.py
+++ b/gimmecert/commands.py
@@ -33,6 +33,7 @@ class ExitCode:
     """
 
     SUCCESS = 0
+    ERROR_ARGUMENTS = 2
     ERROR_ALREADY_INITIALISED = 10
     ERROR_NOT_INITIALISED = 11
     ERROR_CERTIFICATE_ALREADY_ISSUED = 12
@@ -46,7 +47,7 @@ class InvalidCommandInvocation(Exception):
     pass
 
 
-def init(stdout, stderr, project_directory, ca_base_name, ca_hierarchy_depth):
+def init(stdout, stderr, project_directory, ca_base_name, ca_hierarchy_depth, key_generator):
     """
     Initialises the necessary directory and CA hierarchies for use in
     the specified directory.
@@ -66,6 +67,9 @@ def init(stdout, stderr, project_directory, ca_base_name, ca_hierarchy_depth):
     :param ca_hierarchy_depth: Length/depths of CA hierarchy that should be initialised. E.g. total number of CAs in chain.
     :type ca_hierarchy_depth: int
 
+    :param key_generator: Callable for generating private keys.
+    :type key_generator: callable[[], cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey]
+
     :returns: Status code, one from gimmecert.commands.ExitCode.
     :rtype: int
     """
@@ -82,7 +86,7 @@ def init(stdout, stderr, project_directory, ca_base_name, ca_hierarchy_depth):
     gimmecert.storage.initialise_storage(project_directory)
 
     # Generate the CA hierarchy.
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy(ca_base_name, ca_hierarchy_depth)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy(ca_base_name, ca_hierarchy_depth, key_generator)
 
     # Output the CA private keys and certificates.
     for level, (private_key, certificate) in enumerate(ca_hierarchy, 1):
@@ -96,7 +100,7 @@ def init(stdout, stderr, project_directory, ca_base_name, ca_hierarchy_depth):
     full_chain_path = os.path.join(ca_directory, 'chain-full.cert.pem')
     gimmecert.storage.write_certificate_chain(full_chain, full_chain_path)
 
-    print("CA hierarchy initialised. Generated artefacts:", file=stdout)
+    print("CA hierarchy initialised using %s keys. Generated artefacts:" % str(key_generator), file=stdout)
     for level in range(1, ca_hierarchy_depth+1):
         print("    CA Level %d private key: .gimmecert/ca/level%d.key.pem" % (level, level), file=stdout)
         print("    CA Level %d certificate: .gimmecert/ca/level%d.cert.pem" % (level, level), file=stdout)
diff --git a/gimmecert/crypto.py b/gimmecert/crypto.py
index b78cc505b7e1e76967965902dd2ed5c8eb3752e5..d630119bef6cf940ca7c337081e80b23a48e573c 100644
--- a/gimmecert/crypto.py
+++ b/gimmecert/crypto.py
@@ -25,6 +25,73 @@ import cryptography.x509
 from dateutil.relativedelta import relativedelta
 
 
+class KeyGenerator:
+    """
+    Provides abstract factory-like interface for generating private
+    keys. Algorithm and parameters for the private key are provided
+    during instance initialisation by passing-in a specification.
+
+    Instances are callable objects that generate and return the
+    private key according to key specification passed-in during the
+    instance initialisation.
+    """
+
+    def __init__(self, specification):
+        """
+        Initialises an instance.
+
+        :param specification: Specification describing the private keys that that instance should be generating.
+                              For RSA keys, use syntax "rsa:BIT_LENGTH".
+        :type specification: str
+
+        :raises ValueError: If passed-in specification is invalid.
+        """
+
+        try:
+            # This will throw ValueError if we can't get two values
+            # assigned via split.
+            key_type, key_parameters = specification.split(":", 2)
+
+            if key_type == "rsa" and key_parameters.isnumeric():
+                self._algorithm = "rsa"
+                self._parameters = int(key_parameters)
+            else:
+                raise ValueError()
+
+        except ValueError:
+            raise ValueError("Invalid key specification: '%s'" % specification)
+
+    def __str__(self):
+        """
+        Returns string (human-readable) representation of stored key
+        algorithm and parameters.
+
+        :returns: String representation of object.
+        :rtype: str
+        """
+
+        return "%d-bit RSA" % self._parameters
+
+    def __call__(self):
+        """
+        Generates RSA private key. Key size is deterimened by instance's
+        key specification (passed-in during instance creation).
+
+        :returns: RSA private key.
+        :rtype: cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey
+        """
+
+        rsa_public_exponent = 65537
+
+        private_key = cryptography.hazmat.primitives.asymmetric.rsa.generate_private_key(
+            public_exponent=rsa_public_exponent,
+            key_size=self._parameters,
+            backend=cryptography.hazmat.backends.default_backend()
+        )
+
+        return private_key
+
+
 def generate_private_key():
     """
     Generates a 2048-bit RSA private key.
@@ -136,7 +203,7 @@ def issue_certificate(issuer_dn, subject_dn, signing_key, public_key, not_before
     return certificate
 
 
-def generate_ca_hierarchy(base_name, depth):
+def generate_ca_hierarchy(base_name, depth, key_generator):
     """
     Generates CA hierarchy with specified depth, using the provided
     naming as basis for the DNs.
@@ -144,6 +211,9 @@ def generate_ca_hierarchy(base_name, depth):
     :param base_name: Base name for constructing the CA DNs. Resulting DNs are of format 'BASE Level N'.
     :type base_name: str
 
+    :param key_generator: Callable for generating private keys.
+    :type key_generator: callable[[], cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey]
+
     :returns: List of CA private key and certificate pairs, starting with the level 1 (root) CA, and ending with the leaf CA.
     :rtype: list[(cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey, cryptography.x509.Certificate)]
     """
@@ -163,7 +233,7 @@ def generate_ca_hierarchy(base_name, depth):
     for level in range(1, depth+1):
         # Generate info for the new CA.
         dn = get_dn("%s Level %d CA" % (base_name, level))
-        private_key = generate_private_key()
+        private_key = key_generator()
 
         # First certificate issued needs to be self-signed.
         issuer_dn = issuer_dn or dn
diff --git a/tests/conftest.py b/tests/conftest.py
index b848ea7ff7a713000f6e90843715c327b5d19e3d..618868ca1ccb594bdc55284ab2d4fecc9fb21b81 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -86,7 +86,8 @@ def sample_project_directory(tmpdir):
     CSR).
 
     Initialised CA hierarchy is 1 level deep, with basename used being
-    identical to temporary directory base name.
+    identical to temporary directory base name, and it uses 2048-bit
+    RSA keys.
 
     The following server certificates are issued:
 
@@ -142,7 +143,7 @@ def sample_project_directory(tmpdir):
         gimmecert.storage.write_csr(csr, custom_csr_dir.join("%s.csr.pem" % name).strpath)
 
     # Initialise one-level deep hierarchy.
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     # Issue a bunch of certificates.
     for i in range(1, per_type_count + 1):
@@ -170,7 +171,8 @@ def gctmpdir(tmpdir):
     simple CA hierarchy.
 
     Initialised CA hierarchy is 1 level deep, with basename used being
-    identical to temporary directory base name.
+    identical to temporary directory base name, and it uses 2048-bit
+    RSA keys.
 
     The fixture is useful in testing of commands where the CA
     hierarchy does not matter (almost anything except init/status
@@ -184,6 +186,6 @@ def gctmpdir(tmpdir):
     """
 
     # Initialise one-level deep hierarchy.
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     return tmpdir
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 4c6b1a02c6b7771eaffbb6dcf6c8ac28658b5777..44ce1fa5c4cc011d6a5ba6b2ce1d042808a0d825 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -224,6 +224,10 @@ VALID_CLI_INVOCATIONS = [
     ("gimmecert.cli.init", ["gimmecert", "init", "--ca-hierarchy-depth", "3"]),
     ("gimmecert.cli.init", ["gimmecert", "init", "-d", "3"]),
 
+    # init, key specification long and short option
+    ("gimmecert.cli.init", ["gimmecert", "init", "--key-specification", "rsa:4096"]),
+    ("gimmecert.cli.init", ["gimmecert", "init", "-k", "rsa:4096"]),
+
     # server, no options
     ("gimmecert.cli.server", ["gimmecert", "server", "myserver"]),
 
@@ -299,6 +303,52 @@ def test_parser_commands_and_options_are_available(tmpdir, command_function, cli
         gimmecert.cli.main()  # Should not raise
 
 
+# List of _invalid_ CLI invocations to use in
+# test_invalid_parser_commands_and_options_produce_error.
+#
+# Each element in this list should be a tuple where first element is
+# the command function (relative to CLI module) that should be mocked,
+# while second element is list of CLI arguments for invoking the
+# command from CLI. See test documentation for more details.
+INVALID_CLI_INVOCATIONS = [
+    # init, invalid key specification
+    ("gimmecert.cli.init", ["gimmecert", "init", "-k", "rsa"]),
+    ("gimmecert.cli.init", ["gimmecert", "init", "-k", "rsa:not_a_number"]),
+    ("gimmecert.cli.init", ["gimmecert", "init", "-k", "unsupported:algorithm"]),
+]
+
+
+@pytest.mark.parametrize("command_function, cli_invocation", INVALID_CLI_INVOCATIONS)
+def test_invalid_parser_commands_and_options_produce_error(tmpdir, command_function, cli_invocation):
+    """
+    Tests handling of invalid CLI invocations by top-level and command
+    parsers.
+
+    This test helps greatly reduce duplication of code, at the expense
+    of some flexibility.
+
+    The passed-in command_function is mocked and set-up to return a
+    success exit code, since the main point is to ensure the CLI
+    supports specific commands and parameters. E.g. the parser should
+    be the one producing errors.
+
+    To add a new valid invocation of CLI, update the
+    INVALID_CLI_INVOCATIONS variable above.
+    """
+
+    # This should ensure we don't accidentally create artifacts
+    # outside of test directory.
+    tmpdir.chdir()
+
+    with mock.patch(command_function) as mock_command_function, mock.patch('sys.argv', cli_invocation):
+        mock_command_function.return_value = gimmecert.commands.ExitCode.SUCCESS
+
+        with pytest.raises(SystemExit) as e_info:
+            gimmecert.cli.main()
+
+        assert e_info.value.code == gimmecert.commands.ExitCode.ERROR_ARGUMENTS
+
+
 @pytest.mark.parametrize("command", ["help", "init", "server", "client", "renew", "status"])
 @pytest.mark.parametrize("help_option", ["--help", "-h"])
 def test_command_exists_and_accepts_help_flag(tmpdir, command, help_option):
@@ -326,7 +376,8 @@ def test_command_exists_and_accepts_help_flag(tmpdir, command, help_option):
 
 @mock.patch('sys.argv', ['gimmecert', 'init'])
 @mock.patch('gimmecert.cli.init')
-def test_init_command_invoked_with_correct_parameters_no_options(mock_init, tmpdir):
+@mock.patch('gimmecert.cli.KeyGenerator')
+def test_init_command_invoked_with_correct_parameters_no_options(mock_key_generator, mock_init, tmpdir):
     # This should ensure we don't accidentally create artifacts
     # outside of test directory.
     tmpdir.chdir()
@@ -335,14 +386,18 @@ def test_init_command_invoked_with_correct_parameters_no_options(mock_init, tmpd
 
     default_depth = 1
 
+    mock_key_generator.return_value = mock.Mock()
+
     gimmecert.cli.main()
 
-    mock_init.assert_called_once_with(sys.stdout, sys.stderr, tmpdir.strpath, tmpdir.basename, default_depth)
+    mock_key_generator.assert_called_once_with("rsa:2048")
+    mock_init.assert_called_once_with(sys.stdout, sys.stderr, tmpdir.strpath, tmpdir.basename, default_depth, mock_key_generator.return_value)
 
 
 @mock.patch('sys.argv', ['gimmecert', 'init', '-b', 'My Project'])
 @mock.patch('gimmecert.cli.init')
-def test_init_command_invoked_with_correct_parameters_with_options(mock_init, tmpdir):
+@mock.patch('gimmecert.cli.KeyGenerator')
+def test_init_command_invoked_with_correct_parameters_with_options(mock_key_generator, mock_init, tmpdir):
     # This should ensure we don't accidentally create artifacts
     # outside of test directory.
     tmpdir.chdir()
@@ -351,9 +406,11 @@ def test_init_command_invoked_with_correct_parameters_with_options(mock_init, tm
 
     default_depth = 1
 
+    mock_key_generator.return_value = mock.Mock()
+
     gimmecert.cli.main()
 
-    mock_init.assert_called_once_with(sys.stdout, sys.stderr, tmpdir.strpath, 'My Project', default_depth)
+    mock_init.assert_called_once_with(sys.stdout, sys.stderr, tmpdir.strpath, 'My Project', default_depth, mock_key_generator.return_value)
 
 
 @mock.patch('sys.argv', ['gimmecert', 'server'])
diff --git a/tests/test_commands.py b/tests/test_commands.py
index 413f61f74686ceaae6acbf6dab91b2c6f95f3145..737367ca03d4d3f392697a2d389e3d5bac18ea58 100644
--- a/tests/test_commands.py
+++ b/tests/test_commands.py
@@ -26,6 +26,7 @@ import sys
 import cryptography.x509
 
 import gimmecert.commands
+import gimmecert.crypto
 
 import pytest
 from unittest import mock
@@ -37,7 +38,7 @@ def test_init_sets_up_directory_structure(tmpdir):
     ca_dir = tmpdir.join('.gimmecert', 'ca')
     server_dir = tmpdir.join('.gimmecert', 'server')
 
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     assert os.path.exists(base_dir.strpath)
     assert os.path.exists(ca_dir.strpath)
@@ -45,7 +46,7 @@ def test_init_sets_up_directory_structure(tmpdir):
 
 
 def test_init_generates_single_ca_artifact_for_depth_1(tmpdir):
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     assert os.path.exists(tmpdir.join('.gimmecert', 'ca', 'level1.key.pem').strpath)
     assert os.path.exists(tmpdir.join('.gimmecert', 'ca', 'level1.cert.pem').strpath)
@@ -53,7 +54,7 @@ def test_init_generates_single_ca_artifact_for_depth_1(tmpdir):
 
 
 def test_init_generates_three_ca_artifacts_for_depth_3(tmpdir):
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 3)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 3, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     assert os.path.exists(tmpdir.join('.gimmecert', 'ca', 'level1.key.pem').strpath)
     assert os.path.exists(tmpdir.join('.gimmecert', 'ca', 'level1.cert.pem').strpath)
@@ -65,7 +66,7 @@ def test_init_generates_three_ca_artifacts_for_depth_3(tmpdir):
 
 
 def test_init_outputs_full_chain_for_depth_1(tmpdir):
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     level1_certificate = tmpdir.join('.gimmecert', 'ca', 'level1.cert.pem').read()
     full_chain = tmpdir.join('.gimmecert', 'ca', 'chain-full.cert.pem').read()
@@ -74,7 +75,7 @@ def test_init_outputs_full_chain_for_depth_1(tmpdir):
 
 
 def test_init_outputs_full_chain_for_depth_3(tmpdir):
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 3)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 3, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     level1_certificate = tmpdir.join('.gimmecert', 'ca', 'level1.cert.pem').read()
     level2_certificate = tmpdir.join('.gimmecert', 'ca', 'level2.cert.pem').read()
@@ -87,26 +88,26 @@ def test_init_outputs_full_chain_for_depth_3(tmpdir):
 
 
 def test_init_returns_success_if_directory_has_not_been_previously_initialised(tmpdir):
-    status_code = gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
+    status_code = gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     assert status_code == gimmecert.commands.ExitCode.SUCCESS
 
 
 def test_init_returns_error_code_if_directory_has_been_previously_initialised(tmpdir):
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
-    status_code = gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
+    status_code = gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     assert status_code == gimmecert.commands.ExitCode.ERROR_ALREADY_INITIALISED
 
 
 def test_init_does_not_overwrite_artifcats_if_already_initialised(tmpdir):
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     level1_private_key_before = tmpdir.join('.gimmecert', 'ca', 'level1.key.pem').read()
     level1_certificate_before = tmpdir.join('.gimmecert', 'ca', 'level1.cert.pem').read()
     full_chain_before = tmpdir.join('.gimmecert', 'ca', 'chain-full.cert.pem').read()
 
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     level1_private_key_after = tmpdir.join('.gimmecert', 'ca', 'level1.key.pem').read()
     level1_certificate_after = tmpdir.join('.gimmecert', 'ca', 'level1.cert.pem').read()
@@ -249,29 +250,29 @@ def test_init_command_stdout_and_stderr_for_single_ca(tmpdir):
     stdout_stream = io.StringIO()
     stderr_stream = io.StringIO()
 
-    gimmecert.commands.init(stdout_stream, stderr_stream, tmpdir.strpath, "myproject", 1)
+    gimmecert.commands.init(stdout_stream, stderr_stream, tmpdir.strpath, "myproject", 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     stdout = stdout_stream.getvalue()
     stderr = stderr_stream.getvalue()
 
     assert stderr == ""
-    assert "CA hierarchy initialised" in stdout
+    assert "CA hierarchy initialised using 2048-bit RSA keys" in stdout
     assert ".gimmecert/ca/level1.cert.pem" in stdout
     assert ".gimmecert/ca/level1.key.pem" in stdout
     assert ".gimmecert/ca/chain-full.cert.pem" in stdout
 
 
-def test_init_command_stdout_and_stderr_for_multiple_cas(tmpdir):
+def test_init_command_stdout_and_stderr_for_multiple_cas_with_rsa_1024(tmpdir):
     stdout_stream = io.StringIO()
     stderr_stream = io.StringIO()
 
-    gimmecert.commands.init(stdout_stream, stderr_stream, tmpdir.strpath, "myproject", 3)
+    gimmecert.commands.init(stdout_stream, stderr_stream, tmpdir.strpath, "myproject", 3, gimmecert.crypto.KeyGenerator("rsa:1024"))
 
     stdout = stdout_stream.getvalue()
     stderr = stderr_stream.getvalue()
 
     assert stderr == ""
-    assert "CA hierarchy initialised" in stdout
+    assert "CA hierarchy initialised using 1024-bit RSA keys" in stdout
     assert ".gimmecert/ca/level1.cert.pem" in stdout
     assert ".gimmecert/ca/level1.key.pem" in stdout
     assert ".gimmecert/ca/level2.cert.pem" in stdout
@@ -285,9 +286,9 @@ def test_init_command_stdout_and_stderr_if_hierarchy_already_initialised(tmpdir)
     stdout_stream = io.StringIO()
     stderr_stream = io.StringIO()
 
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, "myproject", 1)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, "myproject", 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
-    gimmecert.commands.init(stdout_stream, stderr_stream, tmpdir.strpath, "myproject", 1)
+    gimmecert.commands.init(stdout_stream, stderr_stream, tmpdir.strpath, "myproject", 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     stdout = stdout_stream.getvalue()
     stderr = stderr_stream.getvalue()
@@ -627,7 +628,7 @@ def test_status_reports_ca_hierarchy_information(tmpdir):
     stderr_stream = io.StringIO()
 
     with freeze_time('2018-01-01 00:15:00'):
-        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 3)
+        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 3, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     with freeze_time('2018-06-01 00:15:00'):
         status_code = gimmecert.commands.status(stdout_stream, stderr_stream, tmpdir.strpath)
@@ -676,7 +677,7 @@ def test_status_reports_server_certificate_information(tmpdir):
     gimmecert.storage.write_csr(myserver3_csr, myserver3_csr_file.strpath)
 
     with freeze_time('2018-01-01 00:15:00'):
-        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 3)
+        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 3, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     with freeze_time('2018-02-01 00:15:00'):
         gimmecert.commands.server(io.StringIO(), io.StringIO(), tmpdir.strpath, 'myserver1', None, None)
@@ -744,7 +745,7 @@ def test_status_reports_client_certificate_information(tmpdir):
     gimmecert.storage.write_csr(myclient3_csr, myclient3_csr_file.strpath)
 
     with freeze_time('2018-01-01 00:15:00'):
-        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 3)
+        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 3, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     with freeze_time('2018-02-01 00:15:00'):
         gimmecert.commands.client(io.StringIO(), io.StringIO(), tmpdir.strpath, 'myclient1', None)
@@ -802,7 +803,7 @@ def test_status_reports_no_server_certificates_were_issued(tmpdir):
 
     # Just create some sample data, but no server certificates.
     with freeze_time('2018-01-01 00:15:00'):
-        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
+        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
         gimmecert.commands.client(io.StringIO(), io.StringIO(), tmpdir.strpath, 'myclient1', None)
         gimmecert.commands.client(io.StringIO(), io.StringIO(), tmpdir.strpath, 'myclient2', None)
 
@@ -823,7 +824,7 @@ def test_status_reports_no_client_certificates_were_issued(tmpdir):
 
     # Just create some sample data, but no client certificates.
     with freeze_time('2018-01-01 00:15:00'):
-        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1)
+        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, tmpdir.basename, 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
         gimmecert.commands.server(io.StringIO(), io.StringIO(), tmpdir.strpath, 'myserver1', None, None)
         gimmecert.commands.server(io.StringIO(), io.StringIO(), tmpdir.strpath, 'myserver2', None, None)
 
@@ -862,7 +863,7 @@ def test_certificate_marked_as_not_valid_or_expired_as_appropriate(tmpdir, subje
 
     # Perform action on our fixed issuance date.
     with freeze_time(issuance_date):
-        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, "My Project", 1)
+        gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, "My Project", 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
         gimmecert.commands.server(io.StringIO(), io.StringIO(), tmpdir.strpath, 'myserver', None, None)
         gimmecert.commands.client(io.StringIO(), io.StringIO(), tmpdir.strpath, 'myclient', None)
 
diff --git a/tests/test_crypto.py b/tests/test_crypto.py
index 21c82796cfc5645ff281fd53d6d998ccf505b42f..f5062aaa319ac66d9020c0453c4b3dfb755cf98a 100644
--- a/tests/test_crypto.py
+++ b/tests/test_crypto.py
@@ -27,6 +27,7 @@ from dateutil.relativedelta import relativedelta
 
 import gimmecert.crypto
 
+import pytest
 from freezegun import freeze_time
 
 
@@ -107,7 +108,7 @@ def test_generate_ca_hierarchy_returns_list_with_3_elements_for_depth_3():
     base_name = 'My Project'
     depth = 3
 
-    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth)
+    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     assert isinstance(hierarchy, list)
     assert len(hierarchy) == depth
@@ -117,7 +118,7 @@ def test_generate_ca_hierarchy_returns_list_with_1_element_for_depth_1():
     base_name = 'My Project'
     depth = 1
 
-    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth)
+    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     assert isinstance(hierarchy, list)
     assert len(hierarchy) == depth
@@ -127,7 +128,7 @@ def test_generate_ca_hierarchy_returns_list_of_private_key_certificate_pairs():
     base_name = 'My Project'
     depth = 3
 
-    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth)
+    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     for private_key, certificate in hierarchy:
         assert isinstance(private_key, cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey)
@@ -137,8 +138,9 @@ def test_generate_ca_hierarchy_returns_list_of_private_key_certificate_pairs():
 def test_generate_ca_hierarchy_subject_dns_have_correct_value():
     base_name = 'My Project'
     depth = 3
+    key_generator = gimmecert.crypto.KeyGenerator("rsa:2048")
 
-    level1, level2, level3 = [certificate for _, certificate in gimmecert.crypto.generate_ca_hierarchy(base_name, depth)]
+    level1, level2, level3 = [certificate for _, certificate in gimmecert.crypto.generate_ca_hierarchy(base_name, depth, key_generator)]
 
     assert level1.subject == cryptography.x509.Name(gimmecert.crypto.get_dn('My Project Level 1 CA'))
     assert level2.subject == cryptography.x509.Name(gimmecert.crypto.get_dn('My Project Level 2 CA'))
@@ -149,7 +151,7 @@ def test_generate_ca_hierarchy_issuer_dns_have_correct_value():
     base_name = 'My Project'
     depth = 3
 
-    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth)
+    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     level1_key, level1_certificate = hierarchy[0]
     level2_key, level2_certificate = hierarchy[1]
@@ -164,7 +166,7 @@ def test_generate_ca_hierarchy_private_keys_match_with_public_keys_in_certificat
     base_name = 'My Project'
     depth = 3
 
-    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth)
+    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     level1_private_key, level1_certificate = hierarchy[0]
     level2_private_key, level2_certificate = hierarchy[1]
@@ -179,7 +181,7 @@ def test_generate_ca_hierarchy_cas_have_differing_keys():
     base_name = 'My Project'
     depth = 3
 
-    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth)
+    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     level1_private_key, _ = hierarchy[0]
     level2_private_key, _ = hierarchy[1]
@@ -198,7 +200,7 @@ def test_generate_ca_hierarchy_certificates_have_same_validity():
     base_name = 'My Project'
     depth = 3
 
-    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth)
+    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     _, level1_certificate = hierarchy[0]
     _, level2_certificate = hierarchy[1]
@@ -248,7 +250,7 @@ def test_generate_ca_hierarchy_produces_certificates_with_ca_basic_constraints()
     base_name = 'My test'
     depth = 3
 
-    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth)
+    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     for _, certificate in hierarchy:
         stored_extension = certificate.extensions.get_extension_for_class(cryptography.x509.BasicConstraints)
@@ -261,7 +263,7 @@ def test_generate_ca_hierarchy_produces_certificates_with_ca_basic_constraints()
 
 
 def test_issue_server_certificate_returns_certificate():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -272,7 +274,7 @@ def test_issue_server_certificate_returns_certificate():
 
 
 def test_issue_server_certificate_sets_correct_extensions():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -317,7 +319,7 @@ def test_issue_server_certificate_sets_correct_extensions():
 
 
 def test_issue_server_certificate_has_correct_issuer_and_subject():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 4)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 4, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[3]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -329,7 +331,7 @@ def test_issue_server_certificate_has_correct_issuer_and_subject():
 
 
 def test_issue_server_certificate_has_correct_public_key():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -341,7 +343,7 @@ def test_issue_server_certificate_has_correct_public_key():
 
 @freeze_time('2018-01-01 00:15:00')
 def test_issue_server_certificate_not_before_is_15_minutes_in_past():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -353,7 +355,7 @@ def test_issue_server_certificate_not_before_is_15_minutes_in_past():
 
 def test_issue_server_certificate_not_before_does_not_exceed_ca_validity():
     with freeze_time('2018-01-01 00:15:00'):
-        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
@@ -367,7 +369,7 @@ def test_issue_server_certificate_not_before_does_not_exceed_ca_validity():
 
 def test_issue_server_certificate_not_after_does_not_exceed_ca_validity():
     with freeze_time('2018-01-01 00:15:00'):
-        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
@@ -380,7 +382,7 @@ def test_issue_server_certificate_not_after_does_not_exceed_ca_validity():
 
 
 def test_issue_server_certificate_incorporates_additional_dns_subject_alternative_names():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -401,7 +403,7 @@ def test_issue_server_certificate_incorporates_additional_dns_subject_alternativ
 
 
 def test_issue_client_certificate_returns_certificate():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -412,7 +414,7 @@ def test_issue_client_certificate_returns_certificate():
 
 
 def test_issue_client_certificate_has_correct_issuer_and_subject():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 4)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 4, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[3]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -424,7 +426,7 @@ def test_issue_client_certificate_has_correct_issuer_and_subject():
 
 
 def test_issue_client_certificate_sets_correct_extensions():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -461,7 +463,7 @@ def test_issue_client_certificate_sets_correct_extensions():
 
 
 def test_issue_client_certificate_has_correct_public_key():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -473,7 +475,7 @@ def test_issue_client_certificate_has_correct_public_key():
 
 @freeze_time('2018-01-01 00:15:00')
 def test_issue_client_certificate_not_before_is_15_minutes_in_past():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -485,7 +487,7 @@ def test_issue_client_certificate_not_before_is_15_minutes_in_past():
 
 def test_issue_client_certificate_not_before_does_not_exceed_ca_validity():
     with freeze_time('2018-01-01 00:15:00'):
-        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
@@ -499,7 +501,7 @@ def test_issue_client_certificate_not_before_does_not_exceed_ca_validity():
 
 def test_issue_client_certificate_not_after_does_not_exceed_ca_validity():
     with freeze_time('2018-01-01 00:15:00'):
-        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
@@ -512,7 +514,7 @@ def test_issue_client_certificate_not_after_does_not_exceed_ca_validity():
 
 
 def test_renew_certificate_returns_certificate():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -524,7 +526,7 @@ def test_renew_certificate_returns_certificate():
 
 
 def test_renew_certificate_has_correct_content():
-    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+    ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
     issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
     private_key = gimmecert.crypto.generate_private_key()
@@ -544,7 +546,7 @@ def test_renew_certificate_not_before_is_15_minutes_in_past():
 
     # Initial server certificate.
     with freeze_time('2018-01-01 00:15:00'):
-        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
         issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
         private_key = gimmecert.crypto.generate_private_key()
@@ -561,7 +563,7 @@ def test_renew_certificate_not_before_does_not_exceed_ca_validity():
 
     # Initial server certificate.
     with freeze_time('2018-01-01 00:15:00'):
-        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
         issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
         private_key = gimmecert.crypto.generate_private_key()
@@ -578,7 +580,7 @@ def test_renew_certificate_not_after_does_not_exceed_ca_validity():
 
     # Initial server certificate.
     with freeze_time('2018-01-01 00:15:00'):
-        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1)
+        ca_hierarchy = gimmecert.crypto.generate_ca_hierarchy('My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
         issuer_private_key, issuer_certificate = ca_hierarchy[0]
 
         private_key = gimmecert.crypto.generate_private_key()
@@ -614,3 +616,75 @@ def test_generate_csr_returns_csr_with_passed_in_name():
 
     assert csr.public_key().public_numbers() == private_key.public_key().public_numbers()
     assert csr.subject == expected_subject_dn
+
+
+@pytest.mark.parametrize("key_specification", [
+    "",
+    "rsa",
+    "rsa:not_a_number",
+    "unsupported:algorithm",
+])
+def test_KeyGenerator_raises_exception_for_invalid_specification(key_specification):
+
+    with pytest.raises(ValueError) as e_info:
+        gimmecert.crypto.KeyGenerator(key_specification)
+
+    assert str(e_info.value) == "Invalid key specification: '%s'" % key_specification
+
+
+@pytest.mark.parametrize("key_specification", [
+    "rsa:1024",
+    "rsa:2048",
+    "rsa:4096",
+])
+def test_KeyGenerator_accepts_valid_specifications(key_specification):
+
+    gimmecert.crypto.KeyGenerator(key_specification)  # should not raise
+
+
+def test_KeyGenerator_stores_specification():
+
+    key_generator = gimmecert.crypto.KeyGenerator("rsa:2048")
+
+    assert key_generator._algorithm == "rsa"
+    assert key_generator._parameters == 2048
+
+
+@pytest.mark.parametrize("key_specification, string_representation", [
+    ("rsa:1024", "1024-bit RSA"),
+    ("rsa:2048", "2048-bit RSA"),
+    ("rsa:4096", "4096-bit RSA"),
+])
+def test_KeyGenerator_string_representation(key_specification, string_representation):
+
+    key_generator = gimmecert.crypto.KeyGenerator(key_specification)
+    assert str(key_generator) == string_representation
+
+
+def test_KeyGenerator_instance_returns_rsa_private_key():
+
+    key_generator_1 = gimmecert.crypto.KeyGenerator("rsa:1024")
+    key_generator_2 = gimmecert.crypto.KeyGenerator("rsa:2048")
+
+    private_key_1 = key_generator_1()
+    private_key_2 = key_generator_2()
+
+    assert isinstance(private_key_1, cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey)
+    assert isinstance(private_key_2, cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey)
+
+    assert private_key_1.key_size == 1024
+    assert private_key_2.key_size == 2048
+
+
+@pytest.mark.parametrize("key_generator, expected_bit_size", [
+    (gimmecert.crypto.KeyGenerator("rsa:1024"), 1024),
+    (gimmecert.crypto.KeyGenerator("rsa:2048"), 2048),
+])
+def test_generate_ca_hierarchy_uses_correct_rsa_bit_size(key_generator, expected_bit_size):
+    base_name = "My Test"
+    depth = 3
+
+    hierarchy = gimmecert.crypto.generate_ca_hierarchy(base_name, depth, key_generator)
+
+    for private_key, _ in hierarchy:
+        assert private_key.key_size == expected_bit_size
diff --git a/tests/test_storage.py b/tests/test_storage.py
index 23cf3de5394cd04d09b7efdb5a397cc8e30fdc14..cec4aeb16bb9e8df2d46e26e18dd159f6be9a13d 100644
--- a/tests/test_storage.py
+++ b/tests/test_storage.py
@@ -80,7 +80,7 @@ def test_write_certificate(tmpdir):
 
 def test_write_certificate_chain(tmpdir):
     output_file = tmpdir.join('chain.cert.pem')
-    certificate_chain = [certificate for _, certificate in gimmecert.crypto.generate_ca_hierarchy('My Project', 3)]
+    certificate_chain = [certificate for _, certificate in gimmecert.crypto.generate_ca_hierarchy('My Project', 3, gimmecert.crypto.KeyGenerator("rsa:2048"))]
     level1_pem, level2_pem, level3_pem = [gimmecert.utils.certificate_to_pem(certificate) for certificate in certificate_chain]
 
     gimmecert.storage.write_certificate_chain(certificate_chain, output_file.strpath)
@@ -106,7 +106,7 @@ def test_is_initialised_returns_false_if_directory_is_not_initialised(tmpdir):
 
 def test_read_ca_hierarchy_returns_list_of_ca_private_key_and_certificate_pairs_for_single_ca(tmpdir):
     tmpdir.chdir()
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, 'My Project', 1)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, 'My Project', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     ca_hierarchy = gimmecert.storage.read_ca_hierarchy(tmpdir.join('.gimmecert', 'ca').strpath)
 
@@ -146,7 +146,7 @@ def test_read_certificate_returns_certificate(tmpdir):
 
 def test_read_ca_hierarchy_returns_list_of_ca_private_key_and_certificate_pairs_in_hierarchy_order_for_multiple_cas(tmpdir):
     tmpdir.chdir()
-    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, 'My Project', 4)
+    gimmecert.commands.init(io.StringIO(), io.StringIO(), tmpdir.strpath, 'My Project', 4, gimmecert.crypto.KeyGenerator("rsa:2048"))
 
     ca_hierarchy = gimmecert.storage.read_ca_hierarchy(tmpdir.join('.gimmecert', 'ca').strpath)
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 1fc1f0d4b4d7cebb00291cc87a7930e13a46b98b..76a2b03b2e3e06d11af6465a88eb969be38c1cdb 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -75,7 +75,7 @@ def test_date_range_to_str():
 
 
 def test_get_dns_names_returns_empty_list_if_no_dns_names():
-    issuer_private_key, issuer_certificate = gimmecert.crypto.generate_ca_hierarchy('My Test', 1)[0]
+    issuer_private_key, issuer_certificate = gimmecert.crypto.generate_ca_hierarchy('My Test', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))[0]
     private_key = gimmecert.crypto.generate_private_key()
 
     certificate = gimmecert.crypto.issue_client_certificate(
@@ -91,7 +91,7 @@ def test_get_dns_names_returns_empty_list_if_no_dns_names():
 
 def test_get_dns_names_returns_list_of_dns_names():
 
-    issuer_private_key, issuer_certificate = gimmecert.crypto.generate_ca_hierarchy('My Test', 1)[0]
+    issuer_private_key, issuer_certificate = gimmecert.crypto.generate_ca_hierarchy('My Test', 1, gimmecert.crypto.KeyGenerator("rsa:2048"))[0]
     private_key = gimmecert.crypto.generate_private_key()
 
     certificate = gimmecert.crypto.issue_server_certificate(