Add more precise type annotations for Python 3.8+

[gwerbin]

Hello,

AES.new returns a Union type, which means that a type checker like Mypy cannot know which type exactly is returned. This is a problem if you want to use an attribute that is present on some members of the union but not other members.

For example, CbcMode has an iv attribute, but EcbMode does not. Therefore this code will not type check:

cipher = AES.new(aes_key, AES.MODE_CBC)
print(cipher.iv)

On Python 3.8 and above, we can solve this problem with typing.Literal and typing.overload.

I also replaced Union[bytes, bytearray, memoryview] with ByteString, but if there is a specific reason to use the former I can change it back.

Apologies if I have not followed proper contribution procedures. I didn't see a contributor guide in the repository.

Edit: I see similar issues with several of the other Crypto.Cipher type stubs, e.g. https://github.com/gwerbin/pycryptodome/blob/master/lib/Crypto/Cipher/DES.pyi. I am happy to extend this PR to include the others if the maintainers want to proceed.

[Legrandin]

I am interested in merging this, but I am not sure it works preperly.
For instance, it would be great to detect a wrong construct like:

cipher = AES.new(b'\x00' * 16, AES.MODE_ECB, iv=b'\x00' * 16)

where the developer meant to use AES.MODE_CBC.

The new() function for this case should be declared like this:

@overload
def new(key: Buffer,
        mode: MODE_ECB,
        use_aesni : bool = ...) -> \
        EcbMode: ...

That is, without all parameters that are not used for the ECB mode.

However, even if I do so, I cannot make mypy to flag the error, even after adjusting the type definitions for all the remaining modes.

[Legrandin]

The type definition for new() can be reduced based on which parameter each cipher mode needs.

[Legrandin]

@overload
def new(key: Buffer,
        mode: MODE_ECB,
        use_aesni : bool = ...) -> \
        EcbMode: ...

[Legrandin]

@overload
def new(key: Buffer,
        mode: MODE_CBC,
        iv : Buffer = ...,
        IV : Buffer = ...,
        use_aesni : bool = ...) -> \
        CbcMode: ...

[Legrandin]

@overload
def new(key: Buffer,
        mode: MODE_CFB,
        iv : Buffer = ...,
        IV : Buffer = ...,
        segment_size : int = ...,
        use_aesni : bool = ...) -> \
        CfbMode: ...

[Legrandin]

@overload
def new(key: Buffer,
        mode: MODE_OFB,
        iv : Buffer = ...,
        IV : Buffer = ...,
        use_aesni : bool = ...) -> \
        OfbMode: ...

[Legrandin]

@overload
def new(key: Buffer,
        mode: MODE_CTR,
        nonce : Buffer = ...,
        initial_value : Union[int, Buffer] = ...,
        counter : Dict = ...,
        use_aesni : bool = ...) -> \
        CtrMode: ...

[Legrandin]

@overload
def new(key: Buffer,
        mode: MODE_EAX,
        nonce : Buffer = ...,
        mac_len : int = ...,
        use_aesni : bool = ...) -> \
        EaxMode: ...

[Legrandin]

@overload
def new(key: Buffer,
        mode: MODE_SIV,
        nonce : Buffer = ...,
        use_aesni : bool = ...) -> \
        SivMode: ...

[Legrandin]

@overload
def new(key: Buffer,
        mode: MODE_GCM,
        nonce : Buffer = ...,
        mac_len : int = ...,
        use_aesni : bool = ...) -> \
        GcmMode: ...

[Legrandin]

@overload
def new(key: Buffer,
        mode: MODE_OCB,
        nonce : Buffer = ...,
        mac_len : int = ...,
        use_aesni : bool = ...) -> \
        OcbMode: ...

[Legrandin]

Is AESMode needed?

[gwerbin]

No, it's not needed. But in the absence of an enum, it serves as a kind of grouping that I think helps make the code somewhat self-documenting.

[Autopilot9369]

Any updates on this issue?

[gwerbin]

I just forgot about it @subhamagr.

@Legrandin I admittedly don't know enough about crypto to know which options are needed for each mode. Thanks for the suggestions.

Is there a better way to annotate counter? I assume it's not as general as Dict[Hashable, Any]. What kind of dictionary should be passed to that parameter?

[tayler6000]

@gwerbin Thank you for your work on this issue. I am type annotating my old project right now and stumbled into this issue. I came here to make this very pull request. I'm glad I found your PR before I started working on it. I hope these changes get committed and the PR is accepted soon!

Let me know if you need any help with it!

[Legrandin]

Small reminder for people interested in this feature that the following error will still not be caught by mypy:

cipher = AES.new(b'\x00' * 16, AES.MODE_ECB, iv=b'\x00' * 16)

I tried to tweak the type definition a few times, but with no success...

[Legrandin]

As a note to the future, the biggest problem was that Literal had to imported from typing.

Despite that, the definition:

@overload
def new(key: ByteString,
        mode: ECB_MODE,
        use_aesni : bool = ...) -> \
        EcbMode: ...

leads to the following revealed type:

Overload(def (key: typing.ByteString, mode: MODE_ECB?, use_aesni: builtins.bool =) -> Crypto.Cipher._mode_ecb.EcbMode, 

instead of the desired:

Overload(def (key: typing.ByteString, mode: Literal[1], use_aesni: builtins.bool =) -> Crypto.Cipher._mode_ecb.EcbMode, 

Therefore, the following expression is still not being seen as an error:

cipher = AES.new(b'\x00' * 16, AES.MODE_ECB, iv=b'\x00' * 16)

The only solution seems to be to repeat Literal[1] again:

@overload
def new(key: ByteString,
        mode: Literal[1],
        use_aesni : bool = ...) -> \
        EcbMode: ...

which is not great. Same for the other modes.

[Legrandin]

Merged with b1794ca. Thanks!

[povilasb]

Much appreciated 🙌

When can we expect this to be released with 3.17? :)