Merge pull request #29 from alexander-akhmetov/issue_12

Fix #12 (database encryption key); Fix #8 (stop method)
alexander-akhmetov · Apr 5, 2019 · 1649cdd · 1649cdd
2 parents 0866b1f + 46db8ac
commit 1649cdd
Show file tree

Hide file tree

Showing 9 changed files with 100 additions and 54 deletions.
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
@@ -2,6 +2,14 @@
 Changelog
 =========
 
+[unreleased]
+
+- **Incompatible** default path for files is changed. Now the library uses an md5 hash of the phone number or bot token instead of just a phone number.
+  It should not be noticeable for most cases, but if you rely on locally saved files or database, you need to pass the ``files_directory`` parameter to the ``telegram.client.Telegram``.
+- Fixed problem with randomly raised "Database encryption key is needed" errors during login process. (#12)
+- Fixed `stop` method execution. (#8)
+- Added ``examples/bot_login.py`` example.
+
 [0.8.0] - 2019-03-17
 
 - ``telegram.client.Telegram`` now supports any update type with a new method ``add_update_handler(handler_type, func)``

diff --git a/examples/bot_login.py b/examples/bot_login.py
@@ -0,0 +1,28 @@
+import argparse
+
+from telegram.client import Telegram
+
+
+def bot_get_me(api_id, api_hash, token):
+    tg = Telegram(
+        api_id=api_id,
+        api_hash=api_hash,
+        bot_token=token,
+        database_encryption_key='changeme1234',
+    )
+    # you must call login method before others
+    tg.login()
+
+    result = tg.get_me()
+    result.wait()
+    print(result.update)
+    tg.stop()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('api_id', help='API id')  # https://my.telegram.org/apps
+    parser.add_argument('api_hash', help='API hash')
+    parser.add_argument('token', help='Bot token')
+    args = parser.parse_args()
+    bot_get_me(args.api_id, args.api_hash, args.token)
diff --git a/telegram/client.py b/telegram/client.py
@@ -1,4 +1,5 @@
 import os
+import hashlib
 import time
 import queue
 import signal
@@ -75,7 +76,13 @@ def __init__(
         self._database_encryption_key = database_encryption_key
 
         if not files_directory:
-            files_directory = f'/tmp/.tdlib_files/{self.phone}/'
+            hasher = hashlib.md5()
+            hasher.update(
+                (self.phone or self.bot_token).encode('utf-8')  # type: ignore
+            )
+            directory_name = hasher.hexdigest()
+            files_directory = f'/tmp/.tdlib_files/{directory_name}/'
+
         self.files_directory = files_directory
 
         self._authorized = False
@@ -328,15 +335,22 @@ def idle(self, stop_signals=(signal.SIGINT, signal.SIGTERM, signal.SIGABRT)):
     def _signal_handler(self, signum, frame):
         self._is_enabled = False
 
+    def get_authorization_state(self):
+        logger.debug('Getting authorization state')
+        data = {'@type': 'getAuthorizationState'}
+
+        return self._send_data(data, result_id='getAuthorizationState')
+
     def login(self):
         """
         Login process (blocking)
 
-        Must be called before any other call. It sends initial params to the tdlib, sets database encryption key, etc.
+        Must be called before any other call.
+        It sends initial params to the tdlib, sets database encryption key, etc.
         """
         authorization_state = None
         actions = {
-            None: self._set_initial_params,
+            None: self.get_authorization_state,
             'authorizationStateWaitTdlibParameters': self._set_initial_params,
             'authorizationStateWaitEncryptionKey': self._send_encryption_key,
             'authorizationStateWaitPhoneNumber': self._send_phone_number_or_bot_token,
@@ -355,7 +369,11 @@ def login(self):
 
             if result:
                 result.wait(raise_exc=True)
-                authorization_state = result.update['authorization_state']['@type']
+
+                if result.id == 'getAuthorizationState':
+                    authorization_state = result.update['@type']
+                else:
+                    authorization_state = result.update['authorization_state']['@type']
 
     def _set_initial_params(self) -> AsyncResult:
         logger.info(

diff --git a/telegram/tdjson.py b/telegram/tdjson.py
@@ -111,6 +111,4 @@ def td_execute(self, query: Dict[Any, Any]) -> Dict[Any, Any]:
         return result
 
     def stop(self) -> None:
-        self._tdjson._td_json_client_destroy(
-            self.td_json_client
-        )  # pylint: disable=protected-access
+        self._td_json_client_destroy(self.td_json_client)
diff --git a/telegram/utils.py b/telegram/utils.py
@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
 if TYPE_CHECKING:
-    from telegram.client import Telegram    # noqa  pylint: disable=cyclic-import
+    from telegram.client import Telegram  # noqa  pylint: disable=cyclic-import
 
 
 class AsyncResult:

diff --git a/telegram/worker.py b/telegram/worker.py
@@ -24,7 +24,9 @@ class SimpleWorker(BaseWorker):
     """Simple one-thread worker"""
 
     def run(self) -> None:
-        self._thread = threading.Thread(target=self._run_thread)    # pylint: disable=attribute-defined-outside-init
+        self._thread = threading.Thread(  # pylint: disable=attribute-defined-outside-init
+            target=self._run_thread
+        )
         self._thread.daemon = True
         self._thread.start()
 

diff --git a/tests/test_tdjson.py b/tests/test_tdjson.py
@@ -7,53 +7,47 @@ def test_for_darwin(self, mocker):
         mocked_resource = mocker.Mock()
 
         with mocker.mock_module.patch('telegram.tdjson.platform.system', mocked_system):
-            with mocker.mock_module.patch('telegram.tdjson.pkg_resources.resource_filename',
-                                          mocked_resource):
+            with mocker.mock_module.patch(
+                'telegram.tdjson.pkg_resources.resource_filename', mocked_resource
+            ):
                 _get_tdjson_lib_path()
 
         mocked_resource.assert_called_once_with(
-            'telegram',
-            'lib/darwin/libtdjson.dylib',
+            'telegram', 'lib/darwin/libtdjson.dylib'
         )
 
     def test_for_linux(self, mocker):
         mocked_system = mocker.Mock(return_value='Linux')
         mocked_resource = mocker.Mock(return_value='/tmp/')
 
         with mocker.mock_module.patch('telegram.tdjson.platform.system', mocked_system):
-            with mocker.mock_module.patch('telegram.tdjson.pkg_resources.resource_filename',
-                                          mocked_resource):
+            with mocker.mock_module.patch(
+                'telegram.tdjson.pkg_resources.resource_filename', mocked_resource
+            ):
                 _get_tdjson_lib_path()
 
-        mocked_resource.assert_called_once_with(
-            'telegram',
-            'lib/linux/libtdjson.so',
-        )
+        mocked_resource.assert_called_once_with('telegram', 'lib/linux/libtdjson.so')
 
     def test_for_windows(self, mocker):
         mocked_system = mocker.Mock(return_value='Windows')
         mocked_resource = mocker.Mock(return_value='/tmp/')
 
         with mocker.mock_module.patch('telegram.tdjson.platform.system', mocked_system):
-            with mocker.mock_module.patch('telegram.tdjson.pkg_resources.resource_filename',
-                                          mocked_resource):
+            with mocker.mock_module.patch(
+                'telegram.tdjson.pkg_resources.resource_filename', mocked_resource
+            ):
                 _get_tdjson_lib_path()
 
-        mocked_resource.assert_called_once_with(
-            'telegram',
-            'lib/linux/libtdjson.so',
-        )
+        mocked_resource.assert_called_once_with('telegram', 'lib/linux/libtdjson.so')
 
     def test_unknown(self, mocker):
         mocked_system = mocker.Mock(return_value='Unknown')
         mocked_resource = mocker.Mock(return_value='/tmp/')
 
         with mocker.mock_module.patch('telegram.tdjson.platform.system', mocked_system):
-            with mocker.mock_module.patch('telegram.tdjson.pkg_resources.resource_filename',
-                                          mocked_resource):
+            with mocker.mock_module.patch(
+                'telegram.tdjson.pkg_resources.resource_filename', mocked_resource
+            ):
                 _get_tdjson_lib_path()
 
-        mocked_resource.assert_called_once_with(
-            'telegram',
-            'lib/linux/libtdjson.so',
-        )
+        mocked_resource.assert_called_once_with('telegram', 'lib/linux/libtdjson.so')
diff --git a/tests/test_telegram_methods.py b/tests/test_telegram_methods.py
@@ -236,6 +236,7 @@ def test_get_chat_history(self, telegram):
 
     def test_set_initial_params(self, telegram):
         async_result = telegram._set_initial_params()
+        phone_md5 = '69560384b84c896952ef20352fbce705'
 
         exp_data = {
             '@type': 'setTdlibParameters',
@@ -247,9 +248,9 @@ def test_set_initial_params(self, telegram):
                 'system_version': 'unknown',
                 'application_version': VERSION,
                 'system_language_code': 'en',
-                'database_directory': f'/tmp/.tdlib_files/{PHONE}/database',
+                'database_directory': f'/tmp/.tdlib_files/{phone_md5}/database',
                 'use_message_database': True,
-                'files_directory': f'/tmp/.tdlib_files/{PHONE}/files',
+                'files_directory': f'/tmp/.tdlib_files/{phone_md5}/files',
             },
             '@extra': {'request_id': 'updateAuthorizationState'},
         }
@@ -296,29 +297,35 @@ def test_login_process_should_do_nothing_if_already_authorized(self, telegram):
     def test_login_process_with_phone(self, telegram):
         telegram._authorized = False
 
-        def _get_ar(data):
-            ar = AsyncResult(client=telegram)
+        def _get_async_result(data, request_id=None):
+            result = AsyncResult(client=telegram)
 
-            ar.update = data
+            result.update = data
+            result.id = request_id
 
-            return ar
+            return result
 
         # login process chain
-        telegram._set_initial_params = lambda: _get_ar(
+        telegram.get_authorization_state = lambda: _get_async_result(
+            data={'@type': 'authorizationStateWaitEncryptionKey'},
+            request_id='getAuthorizationState',
+        )
+
+        telegram._set_initial_params = lambda: _get_async_result(
             data={
                 'authorization_state': {'@type': 'authorizationStateWaitEncryptionKey'}
             }
         )
-        telegram._send_encryption_key = lambda: _get_ar(
+        telegram._send_encryption_key = lambda: _get_async_result(
             data={'authorization_state': {'@type': 'authorizationStateWaitPhoneNumber'}}
         )
-        telegram._send_phone_number_or_bot_token = lambda: _get_ar(
+        telegram._send_phone_number_or_bot_token = lambda: _get_async_result(
             data={'authorization_state': {'@type': 'authorizationStateWaitCode'}}
         )
-        telegram._send_telegram_code = lambda: _get_ar(
+        telegram._send_telegram_code = lambda: _get_async_result(
             data={'authorization_state': {'@type': 'authorizationStateWaitPassword'}}
         )
-        telegram._send_password = lambda: _get_ar(
+        telegram._send_password = lambda: _get_async_result(
             data={'authorization_state': {'@type': 'authorizationStateReady'}}
         )
 

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -21,10 +21,7 @@ def test_str(self, mocker):
 
     def test_parse_update_with_error(self):
         ar = AsyncResult(client=None)
-        update = {
-            '@type': 'error',
-            'some': 'data',
-        }
+        update = {'@type': 'error', 'some': 'data'}
 
         assert ar.error is False
         assert ar.error_info is None
@@ -38,10 +35,7 @@ def test_parse_update_with_error(self):
 
     def test_parse_update_ok(self):
         ar = AsyncResult(client=None)
-        update = {
-            '@type': 'ok',
-            'some': 'data',
-        }
+        update = {'@type': 'ok', 'some': 'data'}
 
         ar.parse_update(update)
 
@@ -52,10 +46,7 @@ def test_parse_update_ok(self):
 
     def test_parse_update(self):
         ar = AsyncResult(client=None)
-        update = {
-            '@type': 'some_type',
-            'some': 'data',
-        }
+        update = {'@type': 'some_type', 'some': 'data'}
 
         ar.parse_update(update)