Merge branch 'mr_533'

Legrandin · Dec 9, 2022 · b1794ca · b1794ca
2 parents efb94f7 + 1246a6d
commit b1794ca
Showing 1 changed file with 168 additions and 31 deletions.
diff --git a/lib/Crypto/Cipher/AES.pyi b/lib/Crypto/Cipher/AES.pyi
@@ -1,4 +1,5 @@
-from typing import Union, Tuple, Optional, Dict
+import sys
+from typing import ByteString, Dict, Optional, Tuple, Union, overload, Literal
 
 from Crypto.Cipher._mode_ecb import EcbMode
 from Crypto.Cipher._mode_cbc import CbcMode
@@ -12,36 +13,172 @@ from Crypto.Cipher._mode_gcm import GcmMode
 from Crypto.Cipher._mode_siv import SivMode
 from Crypto.Cipher._mode_ocb import OcbMode
 
-AESMode = int
-
-MODE_ECB: AESMode
-MODE_CBC: AESMode
-MODE_CFB: AESMode
-MODE_OFB: AESMode
-MODE_CTR: AESMode
-MODE_OPENPGP: AESMode
-MODE_CCM: AESMode
-MODE_EAX: AESMode
-MODE_GCM: AESMode
-MODE_SIV: AESMode
-MODE_OCB: AESMode
-
-Buffer = Union[bytes, bytearray, memoryview]
-
-def new(key: Buffer,
-        mode: AESMode,
-        iv : Buffer = ...,
-        IV : Buffer = ...,
-        nonce : Buffer = ...,
-        segment_size : int = ...,
-        mac_len : int = ...,
-        assoc_len : int = ...,
-        initial_value : Union[int, Buffer] = ...,
-        counter : Dict = ...,
-        use_aesni : bool = ...) -> \
-        Union[EcbMode, CbcMode, CfbMode, OfbMode, CtrMode,
-              OpenPgpMode, CcmMode, EaxMode, GcmMode,
-              SivMode, OcbMode]: ...
+if sys.version_info >= (3, 8):
+    MODE_ECB: Literal[1]
+    MODE_CBC: Literal[2]
+    MODE_CFB: Literal[3]
+    MODE_OFB: Literal[5]
+    MODE_CTR: Literal[6]
+    MODE_OPENPGP: Literal[7]
+    MODE_CCM: Literal[8]
+    MODE_EAX: Literal[9]
+    MODE_SIV: Literal[10]
+    MODE_GCM: Literal[11]
+    MODE_OCB: Literal[12]
+
+    # MODE_ECB
+    @overload
+    def new(key: ByteString,
+            mode: Literal[1],
+            use_aesni : bool = ...) -> \
+            EcbMode: ...
+
+    # MODE_CBC
+    @overload
+    def new(key: ByteString,
+            mode: Literal[2],
+            iv : ByteString = ...,
+            use_aesni : bool = ...) -> \
+            CbcMode: ...
+
+    @overload
+    def new(key: ByteString,
+            mode: Literal[2],
+            IV : ByteString = ...,
+            use_aesni : bool = ...) -> \
+            CbcMode: ...
+
+    # MODE_CFB
+    @overload
+    def new(key: ByteString,
+            mode: Literal[3],
+            iv : ByteString = ...,
+            segment_size : int = ...,
+            use_aesni : bool = ...) -> \
+            CfbMode: ...
+
+    @overload
+    def new(key: ByteString,
+            mode: Literal[3],
+            IV : ByteString = ...,
+            segment_size : int = ...,
+            use_aesni : bool = ...) -> \
+            CfbMode: ...
+
+    # MODE_OFB
+    @overload
+    def new(key: ByteString,
+            mode: Literal[5],
+            iv : ByteString = ...,
+            use_aesni : bool = ...) -> \
+            OfbMode: ...
+
+    @overload
+    def new(key: ByteString,
+            mode: Literal[5],
+            IV : ByteString = ...,
+            use_aesni : bool = ...) -> \
+            OfbMode: ...
+
+    # MODE_CTR
+    @overload
+    def new(key: ByteString,
+            mode: Literal[6],
+            nonce : ByteString = ...,
+            initial_value : Union[int, ByteString] = ...,
+            counter : Dict = ...,
+            use_aesni : bool = ...) -> \
+            CtrMode: ...
+
+    # MODE_OPENPGP
+    @overload
+    def new(key: ByteString,
+            mode: Literal[7],
+            iv : ByteString = ...,
+            use_aesni : bool = ...) -> \
+            OpenPgpMode: ...
+
+    @overload
+    def new(key: ByteString,
+            mode: Literal[7],
+            IV : ByteString = ...,
+            use_aesni : bool = ...) -> \
+            OpenPgpMode: ...
+
+    # MODE_CCM
+    @overload
+    def new(key: ByteString,
+            mode: Literal[8],
+            nonce : ByteString = ...,
+            mac_len : int = ...,
+            assoc_len : int = ...,
+            use_aesni : bool = ...) -> \
+            CcmMode: ...
+
+    # MODE_EAX
+    @overload
+    def new(key: ByteString,
+            mode: Literal[9],
+            nonce : ByteString = ...,
+            mac_len : int = ...,
+            use_aesni : bool = ...) -> \
+            EaxMode: ...
+
+    # MODE_GCM
+    @overload
+    def new(key: ByteString,
+            mode: Literal[10],
+            nonce : ByteString = ...,
+            use_aesni : bool = ...) -> \
+            SivMode: ...
+
+    # MODE_SIV
+    @overload
+    def new(key: ByteString,
+            mode: Literal[11],
+            nonce : ByteString = ...,
+            mac_len : int = ...,
+            use_aesni : bool = ...) -> \
+            GcmMode: ...
+
+    # MODE_OCB
+    @overload
+    def new(key: ByteString,
+            mode: Literal[12],
+            nonce : ByteString = ...,
+            mac_len : int = ...,
+            use_aesni : bool = ...) -> \
+            OcbMode: ...
+
+else:
+    AESMode = int
+
+    MODE_ECB: AESMode
+    MODE_CBC: AESMode
+    MODE_CFB: AESMode
+    MODE_OFB: AESMode
+    MODE_CTR: AESMode
+    MODE_OPENPGP: AESMode
+    MODE_CCM: AESMode
+    MODE_EAX: AESMode
+    MODE_GCM: AESMode
+    MODE_SIV: AESMode
+    MODE_OCB: AESMode
+
+    def new(key: ByteString,
+            mode: AESMode,
+            iv : ByteString = ...,
+            IV : ByteString = ...,
+            nonce : ByteString = ...,
+            segment_size : int = ...,
+            mac_len : int = ...,
+            assoc_len : int = ...,
+            initial_value : Union[int, ByteString] = ...,
+            counter : Dict = ...,
+            use_aesni : bool = ...) -> \
+            Union[EcbMode, CbcMode, CfbMode, OfbMode, CtrMode,
+                  OpenPgpMode, CcmMode, EaxMode, GcmMode,
+                  SivMode, OcbMode]: ...
 
 block_size: int
 key_size: Tuple[int, int, int]