Add to_jwk static method to ECAlgorithm (#732)

* Add to_jwk static method to ECAlgorithm

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add in tests for ECAlgorithm.to_jwk

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add to_jwk pull request to changelog

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

4 files changed, 138 insertions, 0 deletions

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 8b3b146..7a4674f 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -15,6 +15,7 @@ Fixed
 
 Added
 ~~~~~
+- Add to_jwk static method to ECAlgorithm by @leonsmith in https://github.com/jpadilla/pyjwt/pull/732
 
 `v2.4.0 <https://github.com/jpadilla/pyjwt/compare/2.3.0...2.4.0>`__
 -----------------------------------------------------------------------
diff --git a/jwt/algorithms.py b/jwt/algorithms.py
index 46a1a53..4c178a3 100644
--- a/jwt/algorithms.py
+++ b/jwt/algorithms.py
@@ -440,6 +440,41 @@ if has_crypto:
                 return False
 
         @staticmethod
+        def to_jwk(key_obj):
+
+            if isinstance(key_obj, EllipticCurvePrivateKey):
+                public_numbers = key_obj.public_key().public_numbers()
+            elif isinstance(key_obj, EllipticCurvePublicKey):
+                public_numbers = key_obj.public_numbers()
+            else:
+                raise InvalidKeyError("Not a public or private key")
+
+            if isinstance(key_obj.curve, ec.SECP256R1):
+                crv = "P-256"
+            elif isinstance(key_obj.curve, ec.SECP384R1):
+                crv = "P-384"
+            elif isinstance(key_obj.curve, ec.SECP521R1):
+                crv = "P-521"
+            elif isinstance(key_obj.curve, ec.SECP256K1):
+                crv = "secp256k1"
+            else:
+                raise InvalidKeyError(f"Invalid curve: {key_obj.curve}")
+
+            obj = {
+                "kty": "EC",
+                "crv": crv,
+                "x": to_base64url_uint(public_numbers.x).decode(),
+                "y": to_base64url_uint(public_numbers.y).decode(),
+            }
+
+            if isinstance(key_obj, EllipticCurvePrivateKey):
+                obj["d"] = to_base64url_uint(
+                    key_obj.private_numbers().private_value
+                ).decode()
+
+            return json.dumps(obj)
+
+        @staticmethod
         def from_jwk(jwk):
             try:
                 if isinstance(jwk, str):
diff --git a/tests/keys/testkey_ec_secp192r1.priv b/tests/keys/testkey_ec_secp192r1.priv
new file mode 100644
index 0000000..0f4d1c7
--- /dev/null
+++ b/tests/keys/testkey_ec_secp192r1.priv
@@ -0,0 +1,5 @@
+-----BEGIN PRIVATE KEY-----
+MG8CAQAwEwYHKoZIzj0CAQYIKoZIzj0DAQEEVTBTAgEBBBiON6kYcPu8ZUDRTu8W
+eXJ2FmX7e9yq0hahNAMyAARHecLjkXWDUJfZ4wiFH61JpmonCYH1GpinVlqw68Sf
+wtDHg2F6SifQEFC6VKj1ZXw=
+-----END PRIVATE KEY-----
diff --git a/tests/test_algorithms.py b/tests/test_algorithms.py
index ac26600..538078a 100644
--- a/tests/test_algorithms.py
+++ b/tests/test_algorithms.py
@@ -237,6 +237,103 @@ class TestAlgorithms:
                 )
 
     @crypto_required
+    def test_ec_private_key_to_jwk_works_with_from_jwk(self):
+        algo = ECAlgorithm(ECAlgorithm.SHA256)
+
+        with open(key_path("testkey_ec.priv")) as ec_key:
+            orig_key = algo.prepare_key(ec_key.read())
+
+        parsed_key = algo.from_jwk(algo.to_jwk(orig_key))
+        assert parsed_key.private_numbers() == orig_key.private_numbers()
+        assert (
+            parsed_key.private_numbers().public_numbers
+            == orig_key.private_numbers().public_numbers
+        )
+
+    @crypto_required
+    def test_ec_public_key_to_jwk_works_with_from_jwk(self):
+        algo = ECAlgorithm(ECAlgorithm.SHA256)
+
+        with open(key_path("testkey_ec.pub")) as ec_key:
+            orig_key = algo.prepare_key(ec_key.read())
+
+        parsed_key = algo.from_jwk(algo.to_jwk(orig_key))
+        assert parsed_key.public_numbers() == orig_key.public_numbers()
+
+    @crypto_required
+    def test_ec_to_jwk_returns_correct_values_for_public_key(self):
+        algo = ECAlgorithm(ECAlgorithm.SHA256)
+
+        with open(key_path("testkey_ec.pub")) as keyfile:
+            pub_key = algo.prepare_key(keyfile.read())
+
+        key = algo.to_jwk(pub_key)
+
+        expected = {
+            "kty": "EC",
+            "crv": "P-256",
+            "x": "HzAcUWSlGBHcuf3y3RiNrWI-pE6-dD2T7fIzg9t6wEc",
+            "y": "t2G02kbWiOqimYfQAfnARdp2CTycsJPhwA8rn1Cn0SQ",
+        }
+
+        assert json.loads(key) == expected
+
+    @crypto_required
+    def test_ec_to_jwk_returns_correct_values_for_private_key(self):
+        algo = ECAlgorithm(ECAlgorithm.SHA256)
+
+        with open(key_path("testkey_ec.priv")) as keyfile:
+            priv_key = algo.prepare_key(keyfile.read())
+
+        key = algo.to_jwk(priv_key)
+
+        expected = {
+            "kty": "EC",
+            "crv": "P-256",
+            "x": "HzAcUWSlGBHcuf3y3RiNrWI-pE6-dD2T7fIzg9t6wEc",
+            "y": "t2G02kbWiOqimYfQAfnARdp2CTycsJPhwA8rn1Cn0SQ",
+            "d": "2nninfu2jMHDwAbn9oERUhRADS6duQaJEadybLaa0YQ",
+        }
+
+        assert json.loads(key) == expected
+
+    @crypto_required
+    def test_ec_to_jwk_raises_exception_on_invalid_key(self):
+        algo = ECAlgorithm(ECAlgorithm.SHA256)
+
+        with pytest.raises(InvalidKeyError):
+            algo.to_jwk({"not": "a valid key"})
+
+    @crypto_required
+    def test_ec_to_jwk_with_valid_curves(self):
+        tests = {
+            "P-256": ECAlgorithm.SHA256,
+            "P-384": ECAlgorithm.SHA384,
+            "P-521": ECAlgorithm.SHA512,
+            "secp256k1": ECAlgorithm.SHA256,
+        }
+        for (curve, hash) in tests.items():
+            algo = ECAlgorithm(hash)
+
+            with open(key_path(f"jwk_ec_pub_{curve}.json")) as keyfile:
+                pub_key = algo.from_jwk(keyfile.read())
+                assert json.loads(algo.to_jwk(pub_key))["crv"] == curve
+
+            with open(key_path(f"jwk_ec_key_{curve}.json")) as keyfile:
+                priv_key = algo.from_jwk(keyfile.read())
+                assert json.loads(algo.to_jwk(priv_key))["crv"] == curve
+
+    @crypto_required
+    def test_ec_to_jwk_with_invalid_curve(self):
+        algo = ECAlgorithm(ECAlgorithm.SHA256)
+
+        with open(key_path("testkey_ec_secp192r1.priv")) as keyfile:
+            priv_key = algo.prepare_key(keyfile.read())
+
+        with pytest.raises(InvalidKeyError):
+            algo.to_jwk(priv_key)
+
+    @crypto_required
     def test_rsa_jwk_public_and_private_keys_should_parse_and_verify(self):
         algo = RSAAlgorithm(RSAAlgorithm.SHA256)