Refine typing stub for AES

lib/Crypto/Cipher/AES.py

@@ -211,9 +211,16 @@ def new(key, mode, *args, **kwargs):
             If not specified, all associated data is buffered internally,
             which may represent a problem for very large messages.
 
-        *   **initial_value** : (*integer*) --
-            (Only ``MODE_CTR``). The initial value for the counter within
-            the counter block. By default it is **0**.
+        *   **initial_value** : (*integer* or *bytes/bytearray/memoryview*) --
+            (Only ``MODE_CTR``).
+            The initial value for the counter. If not present, the cipher will
+            start counting from 0. The value is incremented by one for each block.
+            The counter number is encoded in big endian mode.
+
+        *   **counter** : (*object*) --
+            Instance of ``Crypto.Util.Counter``, which allows full customization
+            of the counter block. This parameter is incompatible to both ``nonce``
+            and ``initial_value``.
 
         *   **use_aesni** : (*boolean*) --
             Use Intel AES-NI hardware extensions (default: use if available).

lib/Crypto/Cipher/AES.pyi

@@ -1,12 +1,16 @@
-from typing import Any, Union, Tuple, Optional
+from typing import Any, Union, Tuple, Optional, Dict, overload
 
+from Crypto.Cipher._mode_ecb import EcbMode
 from Crypto.Cipher._mode_cbc import CbcMode
 from Crypto.Cipher._mode_cfb import CfbMode
-from Crypto.Cipher._mode_ctr import CtrMode
-from Crypto.Cipher._mode_eax import EaxMode
-from Crypto.Cipher._mode_ecb import EcbMode
 from Crypto.Cipher._mode_ofb import OfbMode
+from Crypto.Cipher._mode_ctr import CtrMode
 from Crypto.Cipher._mode_openpgp import OpenPgpMode
+from Crypto.Cipher._mode_ccm import CcmMode
+from Crypto.Cipher._mode_eax import EaxMode
+from Crypto.Cipher._mode_gcm import GcmMode
+from Crypto.Cipher._mode_siv import SivMode
+from Crypto.Cipher._mode_ocb import OcbMode
 
 AESMode = int
 
@@ -18,21 +22,71 @@ MODE_CTR: AESMode
 MODE_OPENPGP: AESMode
 MODE_CCM: AESMode
 MODE_EAX: AESMode
-MODE_SIV: AESMode
 MODE_GCM: AESMode
+MODE_SIV: AESMode
 MODE_OCB: AESMode
 
-def new(__key: Union[bytes, bytearray, memoryview],
+Buffer = Union[bytes, bytearray, memoryview]
+
+@overload
+def new(__key: Buffer,
         __mode: AESMode,
-        iv : Optional[Union[bytes, bytearray, memoryview]],
-        nonce : Optional[Union[bytes, bytearray, memoryview]],
+        use_aesni : Optional[bool]) -> EcbMode: ...
+
+@overload
+def new(__key: Buffer,
+        __mode: AESMode,
+        __iv : Optional[Buffer],
         segment_size : Optional[int],
+        use_aesni : Optional[bool]) -> \
+        Union[CbcMode, CfbMode, OfbMode, OpenPgpMode]: ...
+
+@overload
+def new(__key: Buffer,
+        __mode: AESMode,
+        iv : Optional[Buffer],
+        segment_size : Optional[int],
+        use_aesni : Optional[bool]) -> \
+        Union[CbcMode, CfbMode, OfbMode, OpenPgpMode]: ...
+
+@overload
+def new(__key: Buffer,
+        __mode: AESMode,
+        IV : Optional[Buffer],
+        segment_size : Optional[int],
+        use_aesni : Optional[bool]) -> \
+        Union[CbcMode, CfbMode, OfbMode, OpenPgpMode]: ...
+
+@overload
+def new(__key: Buffer,
+        __mode: AESMode,
+        nonce : Optional[Buffer],
+        initial_value : Optional[Union[int, Buffer]],
+        use_aesni : Optional[bool]) -> CtrMode: ...
+
+@overload
+def new(__key: Buffer,
+        __mode: AESMode,
+	counter : Dict,
+        use_aesni : Optional[bool]) -> CtrMode: ...
+
+@overload
+def new(__key: Buffer,
+        __mode: AESMode,
+        nonce : Optional[Buffer],
         mac_len : Optional[int],
         msg_len : Optional[int],
         assoc_len : Optional[int],
-        initial_value : Optional[int],
         use_aesni : Optional[bool]) -> \
-        Union[CtrMode, CbcMode, EaxMode, EcbMode, OfbMode, CfbMode, OpenPgpMode]: ...
+        Union[CcmMode]: ...
+
+@overload
+def new(__key: Buffer,
+        __mode: AESMode,
+        nonce : Optional[Buffer],
+        mac_len : Optional[int],
+        use_aesni : Optional[bool]) -> \
+        Union[EaxMode, GcmMode, SivMode]: ...
 
 block_size: int
 key_size: Tuple[int, int, int]

lib/Crypto/__init__.pyi

@@ -1,3 +1,4 @@
 from typing import Tuple, Union
 
 version_info : Tuple[int, int, Union[int, str]]
+__version__ : str