From b9dadee6d68592c2780140b1c0198d0eb0a6004a Mon Sep 17 00:00:00 2001
From: Kaushal Rohit <rohit.kg98@gmail.com>
Date: Thu, 15 Oct 2020 17:02:54 +0530
Subject: [PATCH] Tighten pyjwt types using cryptography stubs (#4645)

---
 third_party/3/jwt/algorithms.pyi | 80 +++++++++++++++++++++-----------
 1 file changed, 52 insertions(+), 28 deletions(-)

diff --git a/third_party/3/jwt/algorithms.pyi b/third_party/3/jwt/algorithms.pyi
index 37a7d4067..d7900e3f6 100644
--- a/third_party/3/jwt/algorithms.pyi
+++ b/third_party/3/jwt/algorithms.pyi
@@ -2,6 +2,27 @@ import sys
 from hashlib import _Hash
 from typing import Any, ClassVar, Dict, Generic, Optional, Set, TypeVar, Union
 
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.asymmetric.ec import (
+    EllipticCurvePrivateKey,
+    EllipticCurvePrivateKeyWithSerialization,
+    EllipticCurvePrivateNumbers,
+    EllipticCurvePublicKey,
+    EllipticCurvePublicKeyWithSerialization,
+    EllipticCurvePublicNumbers,
+)
+from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey
+from cryptography.hazmat.primitives.asymmetric.rsa import (
+    RSAPrivateKey,
+    RSAPrivateKeyWithSerialization,
+    RSAPrivateNumbers,
+    RSAPublicKey,
+    RSAPublicKeyWithSerialization,
+    RSAPublicNumbers,
+)
+from cryptography.hazmat.primitives.asymmetric.utils import Prehashed
+from cryptography.hazmat.primitives.hashes import HashAlgorithm
+
 requires_cryptography = Set[str]
 
 def get_default_algorithms() -> Dict[str, Algorithm[Any]]: ...
@@ -33,46 +54,42 @@ class HMACAlgorithm(Algorithm[bytes]):
     SHA384: ClassVar[_HashAlg]
     SHA512: ClassVar[_HashAlg]
     hash_alg: _HashAlg
-    def __init__(self, _HashAlg) -> None: ...
+    def __init__(self, hash_alg: _HashAlg) -> None: ...
     def prepare_key(self, key: Union[str, bytes]) -> bytes: ...
     @staticmethod
     def to_jwk(key_obj: Union[str, bytes]) -> str: ...
     @staticmethod
     def from_jwk(jwk: _LoadsString) -> bytes: ...
 
-# Only defined if cryptography is installed. Types should be tightened when
-# cryptography gets type hints.
-# See https://github.com/python/typeshed/issues/2542
+# Only defined if cryptography is installed.
 class RSAAlgorithm(Algorithm[Any]):
-    SHA256: ClassVar[Any]
-    SHA384: ClassVar[Any]
-    SHA512: ClassVar[Any]
-    hash_alg: Any
-    def __init__(self, hash_alg: Any) -> None: ...
-    def prepare_key(self, key: Any) -> Any: ...
+    SHA256: ClassVar[hashes.SHA256]
+    SHA384: ClassVar[hashes.SHA384]
+    SHA512: ClassVar[hashes.SHA512]
+    hash_alg: Union[HashAlgorithm, Prehashed]
+    def __init__(self, hash_alg: Union[HashAlgorithm, Prehashed]) -> None: ...
+    def prepare_key(self, key: Union[bytes, str, RSAPrivateKey, RSAPublicKey]) -> Union[RSAPrivateKey, RSAPublicKey]: ...
     @staticmethod
-    def to_jwk(key_obj: Any) -> str: ...
-    @staticmethod
-    def from_jwk(jwk: _LoadsString) -> Any: ...
-    def sign(self, msg: bytes, key: Any) -> bytes: ...
-    def verify(self, msg: bytes, key: Any, sig: bytes) -> bool: ...
+    def from_jwk(jwk: Union[_LoadsString, Dict[str, Any]]) -> Union[RSAPrivateKey, RSAPublicKey]: ...
+    def sign(self, msg: bytes, key: RSAPrivateKey) -> bytes: ...
+    def verify(self, msg: bytes, key: RSAPublicKey, sig: bytes) -> bool: ...
 
-# Only defined if cryptography is installed. Types should be tightened when
-# cryptography gets type hints.
-# See https://github.com/python/typeshed/issues/2542
+# Only defined if cryptography is installed.
 class ECAlgorithm(Algorithm[Any]):
-    SHA256: ClassVar[Any]
-    SHA384: ClassVar[Any]
-    SHA512: ClassVar[Any]
-    hash_alg: Any
-    def __init__(self, hash_alg: Any) -> None: ...
-    def prepare_key(self, key: Any) -> Any: ...
+    SHA256: ClassVar[hashes.SHA256]
+    SHA384: ClassVar[hashes.SHA384]
+    SHA512: ClassVar[hashes.SHA512]
+    hash_alg: Union[HashAlgorithm, Prehashed]
+    def __init__(self, hash_alg: Union[HashAlgorithm, Prehashed]) -> None: ...
+    def prepare_key(
+        self, key: Union[bytes, str, EllipticCurvePrivateKey, EllipticCurvePublicKey]
+    ) -> Union[EllipticCurvePrivateKey, EllipticCurvePublicKey]: ...
     @staticmethod
-    def to_jwk(key_obj: Any) -> str: ...
+    def to_jwk(key_obj: Union[EllipticCurvePrivateKeyWithSerialization, EllipticCurvePublicKeyWithSerialization]) -> str: ...
     @staticmethod
-    def from_jwk(jwk: _LoadsString) -> Any: ...
-    def sign(self, msg: bytes, key: Any) -> bytes: ...
-    def verify(self, msg: bytes, key: Any, sig: bytes) -> bool: ...
+    def from_jwk(jwk: _LoadsString) -> Union[EllipticCurvePrivateKey, EllipticCurvePublicKey]: ...
+    def sign(self, msg: bytes, key: EllipticCurvePrivateKey) -> bytes: ...
+    def verify(self, msg: bytes, key: EllipticCurvePublicKey, sig: bytes) -> bool: ...
 
 # Only defined if cryptography is installed. Types should be tightened when
 # cryptography gets type hints.
@@ -80,3 +97,10 @@ class ECAlgorithm(Algorithm[Any]):
 class RSAPSSAlgorithm(RSAAlgorithm):
     def sign(self, msg: bytes, key: Any) -> bytes: ...
     def verify(self, msg: bytes, key: Any, sig: bytes) -> bool: ...
+
+# Only defined if cryptography is installed.
+class Ed25519Algorithm(Algorithm[Any]):
+    def __init__(self, **kwargs: Any) -> None: ...
+    def prepare_key(self, key: Union[str, bytes, Ed25519PrivateKey, Ed25519PublicKey]) -> Any: ...
+    def sign(self, msg: Union[str, bytes], key: Ed25519PrivateKey) -> bytes: ...
+    def verify(self, msg: Union[str, bytes], key: Ed25519PublicKey, sig: Union[str, bytes]) -> bool: ...