More tests

Author: mosquito

README.md

@@ -136,15 +136,15 @@ jwt-rsa keygen [options]
 
 **Options:**
 
-- `-b, --bits`: Number of bits for the RSA key (default: 2048). Choices: 1024, 2048, 4096, 8192.
+- `-b`, `--bits`: Number of bits for the RSA key (default: 2048). Choices: 1024, 2048, 4096, 8192.
 - `--kid`: Key ID. If not provided, one will be generated.
-- `-a, --algorithm`: Algorithm to use (`RS256`, `RS384`, `RS512`). Default: `RS512`.
-- `-u, --use`: Key usage (`sig` for signature, `enc` for encryption). Default: `sig`.
-- `-o, --format`: Output format (`pem`, `jwk`, `base64`). Default: `jwk`.
-- `-r, --raw`: Output raw JSON without indentation.
-- `-k, --save-public`: Path to save the public key.
-- `-K, --save-private`: Path to save the private key.
-- `-f, --force`: Overwrite existing keys if they exist.
+- `-a`, `--algorithm`: Algorithm to use (`RS256`, `RS384`, `RS512`). Default: `RS512`.
+- `-u`, `--use`: Key usage (`sig` for signature, `enc` for encryption). Default: `sig`.
+- `-o`, `--format`: Output format (`pem`, `jwk`, `base64`). Default: `jwk`.
+- `-r`, `--raw`: Output raw JSON without indentation.
+- `-k`, `--save-public`: Path to save the public key.
+- `-K`, `--save-private`: Path to save the private key.
+- `-f`, `--force`: Overwrite existing keys if they exist.
 
 **Examples:**
 
@@ -212,8 +212,8 @@ jwt-rsa testkey -K PRIVATE_KEY_PATH -k PUBLIC_KEY_PATH
 
 **Options:**
 
-- `-K, --private-key`: Path to the private key (required).
-- `-k, --public-key`: Path to the public key (required).
+- `-K`, `--private-key`: Path to the private key (required).
+- `-k`, `--public-key`: Path to the public key (required).
 
 **Examples:**
 
@@ -238,9 +238,9 @@ jwt-rsa pubkey -K PRIVATE_KEY_PATH [options]
 
 **Options:**
 
-- `-K, --private-key`: Path to the private key (required).
-- `-o, --format`: Output format (`pem`, `jwk`, `base64`). Default: `jwk`.
-- `-r, --raw`: Output raw JSON without indentation.
+- `-K`, `--private-key`: Path to the private key (required).
+- `-o`, `--format`: Output format (`pem`, `jwk`, `base64`). Default: `jwk`.
+- `-r`, `--raw`: Output raw JSON without indentation.
 
 **Examples:**
 
@@ -263,11 +263,11 @@ jwt-rsa issue -K PRIVATE_KEY_PATH [options]
 
 **Options:**
 
-- `-K, --private-key`: Path to the private JWT key (required).
+- `-K`, `--private-key`: Path to the private JWT key (required).
 - `--expired`: Token expiration time in seconds (default: `2678400` seconds, which is 31 days).
 - `--nbf`: "Not Before" claim in seconds (default: `-30`).
-- `-I, --no-interactive`: Disable interactive mode. By default, interactive mode is enabled.
-- `-e, --editor`: Editor to use in interactive mode. Defaults to the `EDITOR` environment variable or `vim`.
+- `-I`, `--no-interactive`: Disable interactive mode. By default, interactive mode is enabled.
+- `-e`, `--editor`: Editor to use in interactive mode. Defaults to the `EDITOR` environment variable or `vim`.
 
 **Examples:**
 
@@ -277,7 +277,8 @@ Issue a JWT token with default expiration and interactive mode:
 jwt-rsa issue -K ./private.pem
 ```
 
-By default will be opened the default editor to edit the claims, the format is python dictionary, with comments and pre-filled values:
+By default will be opened the default editor to edit the claims, the format is python dictionary, with comments and
+pre-filled values:
 
 ```python
 # This modules functions and constants are available:
@@ -323,6 +324,8 @@ $ echo '{"foo": "bar"}' | jwt-rsa issue -K /tmp/key -I --expired 3600
 eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.eyJmb28iOiJiYXIiLCJleHAiOjE3Mzg2MzAwNDcsIm5iZiI6MTczNTk1MTYyN30.HRCQ
 ```
 
+In non interactive mode, the input must be a JSON object with the claims to issue the token.
+
 #### `verify`
 
 Parse and verify a JWT token.
@@ -335,10 +338,10 @@ jwt-rsa verify [options] TOKEN
 
 **Options:**
 
-- `-K, --private-key`: Path to the private key.
-- `-k, --public-key`: Path to the public key. If ommited, the public key will be extracted from the private key.
-- `-V, --no-verify`: Do not verify the token's signature.
-- `-I, --no-interactive`: Disable interactive mode. By default, interactive mode is enabled.
+- `-K`, `--private-key`: Path to the private key.
+- `-k`, `--public-key`: Path to the public key. If ommited, the public key will be extracted from the private key.
+- `-V`, `--no-verify`: Do not verify the token's signature.
+- `-I`, `--no-interactive`: Disable interactive mode. By default, interactive mode is enabled.
 
 **Examples:**
 
@@ -369,12 +372,12 @@ jwt-rsa convert PRIVATE_KEY_PATH [options]
 **Options:**
 
 - `private_key`: Path to the source private key (positional argument).
-- `-k, --save-public`: Path to save the converted public key. If omitted,
+- `-k`, `--save-public`: Path to save the converted public key. If omitted,
   the public key will be saved to the same directory as the private key with a `.pub` extension.
-- `-K, --save-private`: Path to save the converted private key.
-- `-o, --format`: Output format (`pem`, `jwk`, `base64`). Default: `jwk`.
-- `-f, --force`: Overwrite existing keys if they exist.
-- `-r, --raw`: Output raw JSON without indentation.
+- `-K`, `--save-private`: Path to save the converted private key.
+- `-o`, `--format`: Output format (`pem`, `jwk`, `base64`). Default: `jwk`.
+- `-f`, `--force`: Overwrite existing keys if they exist.
+- `-r`, `--raw`: Output raw JSON without indentation.
 
 **Examples:**
 

jwt_rsa/rsa.py

@@ -80,13 +80,19 @@ def load_jwk_private_key(jwk: RSAJWKPrivateKey) -> RSAPrivateKey:
     return private_numbers.private_key(default_backend())
 
 
-def load_jwk(jwk_str: str) -> KeyPair:
-    jwk = json.loads(jwk_str)
-    if "d" in jwk:  # Private key
-        private_key = load_jwk_private_key(jwk)
+def load_jwk(jwk: Union[RSAJWKPublicKey, RSAJWKPrivateKey, str]) -> KeyPair:
+    jwk_dict: Union[RSAJWKPublicKey, RSAJWKPrivateKey]
+
+    if isinstance(jwk, str):
+        jwk_dict = json.loads(jwk)
+    else:
+        jwk_dict = jwk
+
+    if "d" in jwk_dict:  # Private key
+        private_key = load_jwk_private_key(jwk_dict)    # type: ignore
         public_key = private_key.public_key()
     else:  # Public key
-        public_key = load_jwk_public_key(jwk)
+        public_key = load_jwk_public_key(jwk_dict)   # type: ignore
         private_key = None
 
     return KeyPair(private=private_key, public=public_key)
@@ -107,8 +113,7 @@ def rsa_to_jwk(
 @overload
 def rsa_to_jwk(    # type: ignore[overload-cannot-match]
     key: RSAPrivateKey, *, kid: str = "", alg: AlgorithmType = "RS256", use: str = "sig",
-) -> RSAJWKPrivateKey:
-    ...
+) -> RSAJWKPrivateKey: ...
 
 
 def rsa_to_jwk(

jwt_rsa/token.py

@@ -34,25 +34,24 @@ class JWT:
 
     def __init__(
         self,
-        private_key: Optional[RSAPrivateKey] = None,
-        public_key: Optional[RSAPublicKey] = None,
-        expires: Optional[int] = None,
+        key: Optional[Union[RSAPrivateKey, RSAPublicKey]],
+        *, expires: Optional[int] = None,
         nbf_delta: Optional[int] = None,
         algorithm: AlgorithmType = "RS512",
         algorithms: Sequence[AlgorithmType] = ALGORITHMS,
         options: Optional[Dict[str, Any]] = None,
     ):
-
         self.__public_key: RSAPublicKey
-        self.__private_key: Optional[RSAPrivateKey] = private_key
-
-        if public_key is None:
-            if isinstance(self.__private_key, RSAPrivateKey):
-                self.__public_key = self.__private_key.public_key()
-            else:
-                raise ValueError("You must provide either a public or private key")
+        self.__private_key: Optional[RSAPrivateKey]
+
+        if isinstance(key, RSAPrivateKey):
+            self.__public_key = key.public_key()
+            self.__private_key = key
+        elif isinstance(key, RSAPublicKey):
+            self.__public_key = key
+            self.__private_key = None
         else:
-            self.__public_key = public_key
+            raise ValueError("You must provide either a public or private key")
 
         self.__jwt = PyJWT(options)
         self.__expires = expires or self.DEFAULT_EXPIRATION

tests/test_utils.py tests/test_cli.pytests/test_utils.py renamed to tests/test_cli.py

@@ -93,6 +93,73 @@ def test_pem_format(capsys):
     )
 
 
+def test_keygen_no_force(capsys, tmp_path):
+    private_path = tmp_path / "private.pem"
+    public_path = tmp_path / "public.pem"
+
+    keygen(
+        parser.parse_args([
+            "keygen", "-o", "pem",
+            "-K", str(private_path), "-k", str(public_path),
+        ])
+    )
+
+    assert private_path.exists()
+    assert public_path.exists()
+
+    public_content = public_path.read_text()
+    private_content = private_path.read_text()
+
+    assert public_content
+    assert private_content
+
+    # Try to generate keys again buy don't overwrite existing
+    keygen(
+        parser.parse_args([
+            "keygen", "-o", "pem",
+            "-K", str(private_path), "-k", str(public_path),
+        ])
+    )
+
+    assert public_content == public_path.read_text()
+    assert private_content == private_path.read_text()
+
+    keygen(
+        parser.parse_args([
+            "keygen", "-o", "pem", "-f",
+            "-K", str(private_path), "-k", str(public_path),
+        ])
+    )
+
+    assert public_content != public_path.read_text()
+    assert private_content != private_path.read_text()
+
+
+def test_keygen_public_key_auto_naming(capsys, tmp_path):
+    private_path = tmp_path / "key"
+    public_path = tmp_path / "key.pub"
+    keygen(parser.parse_args(["keygen", "-o", "pem", "-K", str(private_path)]))
+
+    assert private_path.exists()
+    assert public_path.exists()
+
+    public_content = public_path.read_text()
+    private_content = private_path.read_text()
+
+    assert public_content
+    assert private_content
+
+    private_path.unlink()
+
+    keygen(parser.parse_args(["keygen", "-o", "pem", "-K", str(private_path)]))
+
+    assert private_path.exists()
+    assert public_path.exists()
+
+    assert public_content == public_path.read_text()
+    assert private_content != private_path.read_text()
+
+
 @pytest.mark.skip(reason="TODO")
 def test_rsa_verify(capsys):
     with mock.patch("sys.argv", ["jwt-rsa", "keygen"]):

tests/test_simple.py tests/test_rsa.pytests/test_simple.py renamed to tests/test_rsa.py

@@ -1,5 +1,6 @@
 import base64
 import datetime
+import json
 import os
 import time
 from itertools import product
@@ -10,6 +11,7 @@
 from jwt.exceptions import InvalidAlgorithmError, InvalidSignatureError
 
 from jwt_rsa import rsa
+from jwt_rsa.types import serialization
 from jwt_rsa.token import JWT
 
 
@@ -67,9 +69,7 @@ def test_rsa_sign():
 @pytest.mark.parametrize("expired,nbf", product(expires, nbfs))
 def test_jwt_token(expired, nbf):
     bits = 2048
-    key, public = rsa.generate_rsa(bits)
-
-    jwt = JWT(key, public)
+    jwt = JWT(rsa.generate_rsa(bits).private)
 
     token = jwt.encode(foo="bar", expired=expired, nbf=nbf)
 
@@ -91,17 +91,15 @@ def test_jwt_token(expired, nbf):
 
 def test_jwt_token_invalid_expiration():
     bits = 2048
-    key, public = rsa.generate_rsa(bits)
-
-    jwt = JWT(key, public)
+    jwt = JWT(rsa.generate_rsa(bits).private)
 
     with pytest.raises(ValueError):
         jwt.encode(foo="bar", expired=None, nbf=None)
 
 
 def test_encode_and_decode_with_private_key():
     bits = 2048
-    key, public = rsa.generate_rsa(bits)
+    key, _ = rsa.generate_rsa(bits)
 
     jwt = JWT(key)
     token = jwt.encode(foo="bar")
@@ -115,8 +113,84 @@ def test_decode_only_ability():
 
     token = JWT(key).encode(foo="bar")
 
-    jwt = JWT(None, public)
+    jwt = JWT(public)
     assert "foo" in jwt.decode(token)
 
     with pytest.raises(RuntimeError):
         jwt.encode(foo=None)
+
+
+def test_jwt_init():
+    bits = 2048
+    key, public = rsa.generate_rsa(bits)
+
+    assert JWT(key)
+
+    assert JWT(public)
+
+    with pytest.raises(ValueError):
+        JWT(None)
+
+
+def test_load_jwk():
+    bits = 2048
+    key, public = rsa.generate_rsa(bits)
+
+    jwk = rsa.rsa_to_jwk(key)
+    assert rsa.load_jwk(jwk).private
+    assert rsa.load_jwk(jwk).public
+
+    jwk = json.dumps(rsa.rsa_to_jwk(key))
+    assert rsa.load_jwk(jwk).private
+    assert rsa.load_jwk(jwk).public
+
+    jwk = rsa.rsa_to_jwk(public)
+    assert not rsa.load_jwk(jwk).private
+    assert rsa.load_jwk(jwk).public
+
+    bad_jwk = rsa.rsa_to_jwk(key)
+    bad_jwk["kty"] = "bad"
+
+    with pytest.raises(ValueError):
+        rsa.load_jwk(bad_jwk)
+
+    bad_jwk = rsa.rsa_to_jwk(public)
+    bad_jwk["kty"] = "bad"
+
+    with pytest.raises(ValueError):
+        rsa.load_jwk(bad_jwk)
+
+    with pytest.raises(ValueError):
+        # noinspection PyTypeChecker
+        rsa.rsa_to_jwk(None)    # type: ignore
+
+
+def test_load_public_key(tmp_path):
+    bits = 2048
+    key, public = rsa.generate_rsa(bits)
+
+    public_path = tmp_path / "public.pem"
+    private_path = tmp_path / "private.pem"
+
+    with open(public_path, "wb") as fp:
+        fp.write(public.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo,
+        ))
+
+    with open(private_path, "wb") as fp:
+        fp.write(key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        ))
+
+    assert rsa.load_public_key(public_path)
+    assert rsa.load_public_key(public_path.read_text())
+    assert rsa.load_public_key(rsa.rsa_to_jwk(public))
+    assert rsa.load_public_key(json.dumps(rsa.rsa_to_jwk(public)))
+
+    assert rsa.load_private_key(private_path)
+    assert rsa.load_private_key(private_path.read_text())
+    assert rsa.load_private_key(rsa.rsa_to_jwk(key))
+    assert rsa.load_private_key(json.dumps(rsa.rsa_to_jwk(key)))