Repository: Zapix/mtpylon (language: python)

Issue:
Implement AuthKeyManager and it's Protocol
- [x] Create AuthKey class.
   - [x] Init with `auth_key` auth key is large integer number computed through create authorization process(https://core.telegram.org/mtproto/auth_key)
   - [x] compute `auth_key_hash` as  `sha1(auth_key)`
   - [x] compute `auth_key_aux_hash` as 64 high-order bits of `auth_key_hash`
   - [x] compute `auth_key_id` as  64 low-order bits of `auth_key_hash`
- [x] Create protocol for AuthKeyManager. 
- [ ] Create in memory AuthKeyManager.


Discussion / hints:

Fix patch:
diff --git a/docs/source/3-rsa-key-manager.rst b/docs/source/3-rsa-key-manager.rst
index a661b85..6a77d29 100644
--- a/docs/source/3-rsa-key-manager.rst
+++ b/docs/source/3-rsa-key-manager.rst
@@ -16,6 +16,9 @@ public rsa keys as python `bytes`. Also It returns list of available fingerprint
 
 .. _RsaManager_implementation:
 
+RSA Key Manager implementation
+------------------------------
+
 MTPylon has got simple RSA Key Manager implementation example of usage:
 
 .. code-block:: python
@@ -67,6 +70,10 @@ MTPylon has got simple RSA Key Manager implementation example of usage:
 
 .. _Rsa_Manager_Protocol:
 
+
+RSA Key Manager protocol
+------------------------
+
 You could implement and use your custom rsa manager by implementing protocol
 `mtpylon.crypto.rsa_manager.RasManagerProtocol`
 
diff --git a/docs/source/4-auth-key-and-auth-key-manager.rst b/docs/source/4-auth-key-and-auth-key-manager.rst
new file mode 100644
index 0000000..74e4d44
--- /dev/null
+++ b/docs/source/4-auth-key-and-auth-key-manager.rst
@@ -0,0 +1,129 @@
+AuthKey and AuthKeyManager
+==========================
+
+
+
+.. _auth_key:
+
+AuthKey
+-------
+
+After creating authorization key we will get big `auth_key` number. For more details
+please check `definition here <https://core.telegram.org/mtproto/description#authorization-key-auth-key>`
+that will be used for encrypting/decrypting messages between client and server.
+In mtpylon special class used to simplify work with `auth_key` number `AuthClass`
+
+`AuthKey` stores key number in `value` attribute and has got several properties:
+
+ - `hash` - sha1 hash of auth key as integer
+
+ - `id` - auth key id as 64 lower bits of auth key hash. please check `definition here<https://core.telegram.org/mtproto/description#key-identifier-auth-key-id>`
+
+ - `aux_hash` - aux hash as 64 higher bits of auth key hash please check `definition here<https://core.telegram.org/mtproto/auth_key#dh-key-exchange-complete>`
+
+
+.. _auth_key_manager:
+
+AuthKeyManager
+--------------
+
+AuthKeyManager is a manager that stores authorization keys.
+By default auth keys stores in memory but to manager them you'll use
+coroutines. Methods to work with auth key manager:
+
+- `set_key(self, value: Union[AuthKey, int])` - coroutine to set authorization key
+  accepts auth_key as instance of `AuthKey` class or as 2048-bit interger
+
+- `has_key(self, value: Union[AuthKey, int])` - coroutine to check does
+  manager know about `auth_key` by `auth_key_id` or not. Accepts
+  instance of `AuthKey` class or integer value as `auth_key_id`
+
+- `get_key(self, value: int)` - coroutine to get instance of `AuthKey` by
+  `auth_key_id`. Accepts only integer value. Raises `AuthKeyDoesNotExist`
+  exception if `auth_key_id` not found in manager
+
+- `del_key(self, value: Union[AuthKey, int])` - coroutine to delete information
+  about authorization key. Accepts instance of `AuthKey` or integer as
+  `auth_key_id`. Raises `AuthKeyDoesNotExist`
+  exception if `auth_key_id` not found in manager
+
+
+Example of usage:
+
+
+.. code-block:: python
+
+    from mtpylon.crypto import AuthKey, AuthKeyManager
+
+    auth_key = 0x32234234 # 2048-bit integer value
+    auth_hash = 0x189e4dd9d40b7e7b5c160c4b0313e843b05b5983 # sha1 of auth_key
+    auth_key_id = 0x0313e843b05b5983  # 64 lower bits of auth_hash
+    auth_aux_hash = 0x189e4dd9d40b7e7b  # 64 higher bits of auth_hash
+
+    manager = AuthKeyManager()
+    manager.set_key(auth_key)
+
+    auth_key_data = await manager.get_key(auth_key_id)
+
+    assert auth_key_data.id == auth_key_id
+    assert auth_key_data.value == auth_key
+    assert auth_key_data.hash == auth_hash
+    assert auth_key_data.aux_hash == auth_aux_hash
+
+
+
+
+.. _auth_key_manger_protocol:
+
+AuthKeyManagerProtocol
+----------------------
+
+Probably you want to store auth_key not in memory but in redis store or other
+key value store. To create you own auth manager you should just implements
+their protocol:
+
+
+.. code-block:: python
+
+
+    class AuthKeyManagerProtocol(Protocol):
+
+        async def set_key(self, value: Union[AuthKey, int]):  # pragma: nocover
+            """
+            Sets new available authentication key. This method
+            takes instance of AuthKey class or auth_key as integer
+            It'll be async function coz we could use separated service to store
+            authorization keys
+            """
+            ...
+
+        async def has_key(
+                self,
+                value: Union[AuthKey, int]
+        ) -> bool:  # pragma: nocover
+            """
+            Checks is this AuthKey store in manager or not. This method
+            takes intance of AuthKey class or integer as auth_key_id.
+            It'll be async function coz we could use separated service to store
+            authorization keys
+            """
+            ...
+
+        async def get_key(self, value: int) -> AuthKey:  # pragma: nocover
+            """
+            Gets auth key data by its id.
+            It'll be async function coz we could separated service to store
+            authorization keys
+
+            Raises: AuthKeyDoesNotExist if key doesn't store in manager
+            """
+            ...
+
+        async def del_key(self, value: Union[AuthKey, int]):
+            """
+            Delete auth key from manager by it's id, value or AuthKey instance
+            It'll be async function coz we could use separated service to store
+            authorization keys
+
+            Raises:  AuthKeyDoesNotExist if key doesn't store in manager
+            """
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 356a282..85ebd1b 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -12,6 +12,8 @@ Welcome to mtpylon's documentation!
 
    1-combinators-and-constructors
    2-schema-and-its-serializers
+   3-rsa-key-manager
+   4-auth-key-and-auth-key-manager
 
 Indices and tables
 ==================
diff --git a/mtpylon/crypto/__init__.py b/mtpylon/crypto/__init__.py
index b1b5c31..1a58d88 100644
--- a/mtpylon/crypto/__init__.py
+++ b/mtpylon/crypto/__init__.py
@@ -1,4 +1,13 @@
 # -*- coding: utf-8 -*-
+from .auth_key import AuthKey
+from .auth_key_manager import AuthKeyManager
+from .exceptions import AuthKeyDoesNotExist
 from .rsa_manager import RsaManager, KeyPair
 
-__all__ = ['RsaManager', 'KeyPair']
+__all__ = [
+    'RsaManager',
+    'KeyPair',
+    'AuthKey',
+    'AuthKeyManager',
+    'AuthKeyDoesNotExist',
+]
diff --git a/mtpylon/crypto/auth_key.py b/mtpylon/crypto/auth_key.py
new file mode 100644
index 0000000..672fa73
--- /dev/null
+++ b/mtpylon/crypto/auth_key.py
@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+from dataclasses import dataclass
+from hashlib import sha1
+
+
+@dataclass()
+class AuthKey:
+    value: int
+
+    def __hash__(self) -> int:
+        """
+        hash truncates to Py_ssize_t value
+        """
+        return self.hash
+
+    @property
+    def hash(self) -> int:
+        return int.from_bytes(
+            sha1(self.value.to_bytes(256, 'big')).digest(),
+            'big'
+        )
+
+    @property
+    def id(self) -> int:
+        return int.from_bytes(
+            self.hash.to_bytes(20, 'big')[-8:],
+            'big'
+        )
+
+    @property
+    def aux_hash(self) -> int:
+        return int.from_bytes(
+            self.hash.to_bytes(20, 'big')[:8],
+            'big'
+        )
diff --git a/mtpylon/crypto/auth_key_manager.py b/mtpylon/crypto/auth_key_manager.py
new file mode 100644
index 0000000..bf9a571
--- /dev/null
+++ b/mtpylon/crypto/auth_key_manager.py
@@ -0,0 +1,93 @@
+# -*- coding: utf-8 -*-
+import asyncio
+from typing import Protocol, Union, Dict
+
+from .auth_key import AuthKey
+from .exceptions import AuthKeyDoesNotExist
+
+
+class AuthKeyManagerProtocol(Protocol):
+
+    async def set_key(self, value: Union[AuthKey, int]):  # pragma: nocover
+        """
+        Sets new available authentication key. This method
+        takes instance of AuthKey class or auth_key as integer
+        It'll be async function coz we could use separated service to store
+        authorization keys
+        """
+        ...
+
+    async def has_key(
+            self,
+            value: Union[AuthKey, int]
+    ) -> bool:  # pragma: nocover
+        """
+        Checks is this AuthKey store in manager or not. This method
+        takes intance of AuthKey class or integer as auth_key_id.
+        It'll be async function coz we could use separated service to store
+        authorization keys
+        """
+        ...
+
+    async def get_key(self, value: int) -> AuthKey:  # pragma: nocover
+        """
+        Gets auth key data by its id.
+        It'll be async function coz we could separated service to store
+        authorization keys
+
+        Raises: AuthKeyDoesNotExist if key doesn't store in manager
+        """
+        ...
+
+    async def del_key(self, value: Union[AuthKey, int]):  # pragma: nocover
+        """
+        Delete auth key from manager by it's id, value or AuthKey instance
+        It'll be async function coz we could use separated service to store
+        authorization keys
+
+        Raises:  AuthKeyDoesNotExist if key doesn't store in manager
+        """
+        ...
+
+
+class AuthKeyManager(AuthKeyManagerProtocol):
+
+    def __init__(self):
+        self._manager_lock = asyncio.Lock()
+        self._manager_map: Dict[int, AuthKey] = {}
+
+    async def set_key(self, value: Union[AuthKey, int]):
+        if isinstance(value, int):
+            value = AuthKey(value)
+
+        async with self._manager_lock:
+            self._manager_map[value.id] = value
+
+    async def has_key(self, value: Union[AuthKey, int]) -> bool:
+        if isinstance(value, AuthKey):
+            value = value.id
+
+        async with self._manager_lock:
+            result = value in self._manager_map
+
+        return result
+
+    async def get_key(self, value: int) -> AuthKey:
+        async with self._manager_lock:
+            try:
+                result = self._manager_map[value]
+            except KeyError:
+                raise AuthKeyDoesNotExist(
+                    f'Auth key id {value} does not exist'
+                )
+        return result
+
+    async def del_key(self, value: Union[AuthKey, int]):
+        if isinstance(value, AuthKey):
+            value = value.id
+
+        async with self._manager_lock:
+            try:
+                del self._manager_map[value]
+            except KeyError:
+                raise AuthKeyDoesNotExist
diff --git a/mtpylon/crypto/exceptions.py b/mtpylon/crypto/exceptions.py
new file mode 100644
index 0000000..e74a1f7
--- /dev/null
+++ b/mtpylon/crypto/exceptions.py
@@ -0,0 +1,6 @@
+# -*- coding: utf-8 -*-
+
+class AuthKeyDoesNotExist(Exception):
+    """
+    Raises when we couldn't find auth key in manager
+    """
diff --git a/setup.cfg b/setup.cfg
index 601cfea..3505bf0 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -2,4 +2,5 @@
 norecursedirs = .svn docs _build tmp*
 
 [flake8]
+ignore = W601
 exclude = docs/


Test patch:
diff --git a/tests/test_crypto/test_auth_key.py b/tests/test_crypto/test_auth_key.py
new file mode 100644
index 0000000..dec9669
--- /dev/null
+++ b/tests/test_crypto/test_auth_key.py
@@ -0,0 +1,37 @@
+# -*- coding: utf-8 -*-
+import pytest
+
+from mtpylon.crypto import AuthKey
+
+auth_key_tests = [
+    pytest.param(
+        0x4fa4120d3646e1cadc7e21fcc9c46111ff8467665908a56b18bd38bee60ee1cccc2eff69dda5e638be2b06e813e6d9832142a054d22f405d9d416f79168140d373b048f55924b836ec5be2ab92e29624edf67bd5d763f8a15c3aed587e7ac70f8fe0d78de45c4b9ea5b81a1e0a098eb064dc9609fa37ca33f703b48461aed343aabc1498dd4d1cbbf3ca4c56cc8c7dc315e251e28312ed258c74326729118251f9153f7ac9d52bd47fdf7072963f6330f06bb72e2cc744af91df3302800e80c23c25a402b97bfc5292589b1c688bd1fde0e9997d9996a32ebd39ba258137b123fa762fd07548c44fd1b9321778f6ec3464ae39402c0445bd02fa8223b30cdfc8,  # noqa
+        0x189e4dd9d40b7e7b5c160c4b0313e843b05b5983,
+        0x0313e843b05b5983,
+        0x189e4dd9d40b7e7b,
+    ),
+    pytest.param(
+        0x73ec5d2b95d59d33a1760dbe97870b617500cd1d05dabe2dca82306e85392e8b78f250507e1391730f6056d1bf7bcadb57bb0c34c8766a0d449e195e93e906ff8c70d77a00880939b7cea5a1bb2659d5e67d58cf8e2dd353ddac7b2b608eaeafd9fe3c44b24ba63c51bcbd50b71624bf8a25a3474fb5e848b84c143f7564cc6362ec2e83e9cd25f7b78361fe757d726ed19fbe569f91def4116ed0abafe09a6992c8d9ceb57c8527c49a1f81b0d087d96aa86dab95182d6cfd09c15f1404642b7ed61b63544918f8566044f53ce1e45fc20bad61ccf6da10b9f9d544c83f3bbaeaadbefcbee74076aeb94e3fe65b711bc727fc32dd03df5e10cbe7ceecb294ce,  # noqa
+        0xce10fe6e8e29cffb284ed5662f8ba7d92e35cfa3,
+        0x2f8ba7d92e35cfa3,
+        0xce10fe6e8e29cffb,
+    ),
+    pytest.param(
+        0x7190f5bcccacee71551570df063b6a4abe1bf603fd079c62351115051bfbeb00315266f9e7963970ecdb893b75f4de90ecadc2e6770cffce315d27ad58ffe9255b07373565bb3dc57f0bf398d8fd04469105faa24a8bc9498e9de4840c322a52d4a88613b2902ba031ee9b9b609912db13a459252ca9249909207193c829142bb7fdb6dd1427461d17a68541ced196763b6b3a72cf9e75b1e60405f53c1a5816adb35aa09ce334490dfb4c05a9801005b51bfdbbd52bb3a2d24d362351494d3cee71213aeda8d06c92ca2e206ed525bec6c99f53aaca6d2fff65c6a627cdb069ac4d769811f1f7d09587a718d2e472def8ce4fdcb7342531a30e894fff0f2a53,  # noqa
+        0x03aa5f90a079852b27281e57fb8a0a8d4da8d551,
+        0xfb8a0a8d4da8d551,
+        0x03aa5f90a079852b,
+    )
+]
+
+
+@pytest.mark.parametrize(
+    'key_value,auth_key_hash,auth_key_id,auth_key_aux_hash',
+    auth_key_tests
+)
+def test_auth_key(key_value, auth_key_hash, auth_key_id, auth_key_aux_hash):
+    auth_key = AuthKey(key_value)
+
+    assert auth_key.hash == auth_key_hash
+    assert auth_key.id == auth_key_id
+    assert auth_key.aux_hash == auth_key_aux_hash
diff --git a/tests/test_crypto/test_auth_key_manager.py b/tests/test_crypto/test_auth_key_manager.py
new file mode 100644
index 0000000..2d27641
--- /dev/null
+++ b/tests/test_crypto/test_auth_key_manager.py
@@ -0,0 +1,135 @@
+# -*- coding: utf-8 -*-
+import pytest
+
+
+from mtpylon.crypto import AuthKey, AuthKeyManager, AuthKeyDoesNotExist
+
+auth_key_tests = [
+    (
+        0x4fa4120d3646e1cadc7e21fcc9c46111ff8467665908a56b18bd38bee60ee1cccc2eff69dda5e638be2b06e813e6d9832142a054d22f405d9d416f79168140d373b048f55924b836ec5be2ab92e29624edf67bd5d763f8a15c3aed587e7ac70f8fe0d78de45c4b9ea5b81a1e0a098eb064dc9609fa37ca33f703b48461aed343aabc1498dd4d1cbbf3ca4c56cc8c7dc315e251e28312ed258c74326729118251f9153f7ac9d52bd47fdf7072963f6330f06bb72e2cc744af91df3302800e80c23c25a402b97bfc5292589b1c688bd1fde0e9997d9996a32ebd39ba258137b123fa762fd07548c44fd1b9321778f6ec3464ae39402c0445bd02fa8223b30cdfc8,  # noqa
+        0x189e4dd9d40b7e7b5c160c4b0313e843b05b5983,
+        0x0313e843b05b5983,
+        0x189e4dd9d40b7e7b,
+    ),
+    (
+        0x73ec5d2b95d59d33a1760dbe97870b617500cd1d05dabe2dca82306e85392e8b78f250507e1391730f6056d1bf7bcadb57bb0c34c8766a0d449e195e93e906ff8c70d77a00880939b7cea5a1bb2659d5e67d58cf8e2dd353ddac7b2b608eaeafd9fe3c44b24ba63c51bcbd50b71624bf8a25a3474fb5e848b84c143f7564cc6362ec2e83e9cd25f7b78361fe757d726ed19fbe569f91def4116ed0abafe09a6992c8d9ceb57c8527c49a1f81b0d087d96aa86dab95182d6cfd09c15f1404642b7ed61b63544918f8566044f53ce1e45fc20bad61ccf6da10b9f9d544c83f3bbaeaadbefcbee74076aeb94e3fe65b711bc727fc32dd03df5e10cbe7ceecb294ce,  # noqa
+        0xce10fe6e8e29cffb284ed5662f8ba7d92e35cfa3,
+        0x2f8ba7d92e35cfa3,
+        0xce10fe6e8e29cffb,
+    ),
+    (
+        0x7190f5bcccacee71551570df063b6a4abe1bf603fd079c62351115051bfbeb00315266f9e7963970ecdb893b75f4de90ecadc2e6770cffce315d27ad58ffe9255b07373565bb3dc57f0bf398d8fd04469105faa24a8bc9498e9de4840c322a52d4a88613b2902ba031ee9b9b609912db13a459252ca9249909207193c829142bb7fdb6dd1427461d17a68541ced196763b6b3a72cf9e75b1e60405f53c1a5816adb35aa09ce334490dfb4c05a9801005b51bfdbbd52bb3a2d24d362351494d3cee71213aeda8d06c92ca2e206ed525bec6c99f53aaca6d2fff65c6a627cdb069ac4d769811f1f7d09587a718d2e472def8ce4fdcb7342531a30e894fff0f2a53,  # noqa
+        0x03aa5f90a079852b27281e57fb8a0a8d4da8d551,
+        0xfb8a0a8d4da8d551,
+        0x03aa5f90a079852b,
+    )
+]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'auth_key,auth_hash,auth_key_id,auth_aux_hash',
+    auth_key_tests
+)
+async def test_set_auth_key(auth_key, auth_hash, auth_key_id, auth_aux_hash):
+    manager = AuthKeyManager()
+
+    assert not await manager.has_key(auth_key_id)
+
+    await manager.set_key(auth_key)
+
+    assert await manager.has_key(auth_key_id)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'auth_key',
+    [item[0] for item in auth_key_tests]
+)
+async def test_has_auth_key_data(auth_key):
+    manager = AuthKeyManager()
+
+    auth_key_data = AuthKey(auth_key)
+    await manager.set_key(auth_key_data)
+
+    assert await manager.has_key(auth_key_data)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'auth_key_id',
+    [
+        0x0313e843b05b5983,
+        0x2f8ba7d92e35cfa3,
+        0xfb8a0a8d4da8d551,
+    ]
+)
+async def test_auth_key_get_error(auth_key_id):
+    manager = AuthKeyManager()
+
+    with pytest.raises(AuthKeyDoesNotExist):
+        await manager.get_key(auth_key_id)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'auth_key,auth_hash,auth_key_id,auth_aux_hash',
+    auth_key_tests
+)
+async def test_auth_key_get(auth_key, auth_hash, auth_key_id, auth_aux_hash):
+    manager = AuthKeyManager()
+    await manager.set_key(auth_key)
+
+    auth_key_data = await manager.get_key(auth_key_id)
+
+    assert auth_key_data.id == auth_key_id
+    assert auth_key_data.value == auth_key
+    assert auth_key_data.hash == auth_hash
+    assert auth_key_data.aux_hash == auth_aux_hash
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'auth_key_id',
+    [
+        0x0313e843b05b5983,
+        0x2f8ba7d92e35cfa3,
+        0xfb8a0a8d4da8d551,
+    ]
+)
+async def test_auth_key_del_error(auth_key_id):
+    manager = AuthKeyManager()
+
+    with pytest.raises(AuthKeyDoesNotExist):
+        await manager.del_key(auth_key_id)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'auth_key,auth_key_id',
+    [(item[0], item[2]) for item in auth_key_tests]
+)
+async def test_auth_key_del_by_id(auth_key, auth_key_id):
+    manager = AuthKeyManager()
+
+    await manager.set_key(auth_key)
+    assert await manager.has_key(auth_key_id)
+    await manager.del_key(auth_key_id)
+
+    assert not await manager.has_key(auth_key_id)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'auth_key,auth_key_id',
+    [(item[0], item[2]) for item in auth_key_tests]
+)
+async def test_auth_key_del_by_auth_data(auth_key, auth_key_id):
+    manager = AuthKeyManager()
+
+    await manager.set_key(auth_key)
+
+    auth_key_data = await manager.get_key(auth_key_id)
+    await manager.del_key(auth_key_data)
+
+    assert not await manager.has_key(auth_key_data)