From 2046215e37076e95a63f2bcf63c3024be348bc82 Mon Sep 17 00:00:00 2001
From: Kishi85 <git@eightyfive.net>
Date: Thu, 4 Apr 2019 10:30:44 +0200
Subject: [PATCH] tools: encapsulate key parameter determination in tools
 function

This is will isolate more cryptographic functions in the tools module
and allow for easier replacement of any cryptographic function should
the need ever arise
---
 acertmgr/authority/v1.py | 10 +++-------
 acertmgr/authority/v2.py | 10 ++--------
 acertmgr/tools.py        | 23 +++++++++++++++++++----
 3 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/acertmgr/authority/v1.py b/acertmgr/authority/v1.py
index 06579b8..9dbba2a 100644
--- a/acertmgr/authority/v1.py
+++ b/acertmgr/authority/v1.py
@@ -29,14 +29,10 @@ class ACMEAuthority(AbstractACMEAuthority):
     # @param key the account key
     # @return the header for ACME
     def _prepare_header(self):
-        numbers = self.key.public_key().public_numbers()
+        alg, jwk = tools.get_key_alg_and_jwk(self.key)
         header = {
-            "alg": "RS256",
-            "jwk": {
-                "e": tools.bytes_to_base64url(tools.number_to_byte_format(numbers.e)),
-                "kty": "RSA",
-                "n": tools.bytes_to_base64url(tools.number_to_byte_format(numbers.n)),
-            },
+            "alg": alg,
+            "jwk": jwk,
         }
         return header
 
diff --git a/acertmgr/authority/v2.py b/acertmgr/authority/v2.py
index dee3a4b..f19e1ed 100644
--- a/acertmgr/authority/v2.py
+++ b/acertmgr/authority/v2.py
@@ -44,16 +44,10 @@ class ACMEAuthority(AbstractACMEAuthority):
             print("API directory retrieval failed ({}). Guessed necessary values: {}".format(code, self.directory))
         self.nonce = None
 
-        # @todo: Add support for key-types other than RSA
-        numbers = key.public_key().public_numbers()
-        self.algorithm = "RS256"
+        self.algorithm, jwk = tools.get_key_alg_and_jwk(key)
         self.account_protected = {
             "alg": self.algorithm,
-            "jwk": {
-                "kty": "RSA",
-                "e": tools.bytes_to_base64url(tools.number_to_byte_format(numbers.e)),
-                "n": tools.bytes_to_base64url(tools.number_to_byte_format(numbers.n)),
-            },
+            "jwk": jwk
         }
         self.account_id = None  # will be updated to correct value during account registration
 
diff --git a/acertmgr/tools.py b/acertmgr/tools.py
index aecb8a7..a2db6fb 100644
--- a/acertmgr/tools.py
+++ b/acertmgr/tools.py
@@ -191,12 +191,27 @@ def convert_der_bytes_to_cert(data):
     return x509.load_der_x509_certificate(data, default_backend())
 
 
+# @brief determine key signing algorithm and jwk data
+# @return key algorithm, signature algorithm, key numbers as a dict
+def get_key_alg_and_jwk(key):
+    if isinstance(key, rsa.RSAPrivateKey):
+        # See https://tools.ietf.org/html/rfc7518#section-6.3
+        numbers = key.public_key().public_numbers()
+        return "RS256", {"kty": "RSA",
+                         "e": bytes_to_base64url(number_to_byte_format(numbers.e)),
+                         "n": bytes_to_base64url(number_to_byte_format(numbers.n))}
+    else:
+        raise ValueError("Unsupported key: {}".format(key))
+
+
 # @brief sign string with key
 def signature_of_str(key, string):
-    # @todo check why this padding is not working
-    # pad = padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH)
-    pad = padding.PKCS1v15()
-    return key.sign(string.encode('utf8'), pad, hashes.SHA256())
+    alg, _ = get_key_alg_and_jwk(key)
+    data = string.encode('utf8')
+    if alg == 'RS256':
+        return key.sign(data, padding.PKCS1v15(), hashes.SHA256())
+    else:
+        raise ValueError("Unsupported signature algorithm: {}".format(alg))
 
 
 # @brief hash a string