Шифрование токенов

.gitignore

@@ -6,3 +6,4 @@ __pycache__
 *.pyc
 config.json
 docker-compose-release.yaml
+docs/build

docker-entrypoint.sh

@@ -11,4 +11,5 @@ fi
 
 sleep 10
 aerich upgrade
+python migrate.py
 python main.py

example.env

@@ -5,6 +5,8 @@ POSTGRES_PASSWORD=SOME_RANDOM_PASSWORD_HERE  # example: x2y0n27ihiez93kmzj82
 POSTGRES_DB=olgram
 POSTGRES_HOST=postgres
 
+TOKEN_ENCRYPTION_KEY=SOME_RANDOM_KEY  # example:
+
 WEBHOOK_HOST=YOUR_HOST_HERE  # example: 11.143.142.140 or my_domain.com
 WEBHOOK_PORT=8443  # allowed: 80, 443, 8080, 8443
 CUSTOM_CERT=true  # use that if you don't set up your own domain and let's encrypt certificate

main.py

@@ -47,7 +47,7 @@ def main():
     loop.run_until_complete(initialization())
 
     loop.create_task(dp.start_polling())
-    loop.create_task(server_main().start())
+    # loop.create_task(server_main().start())
 
     loop.run_forever()
 

migrate.py

@@ -0,0 +1,9 @@
+import asyncio
+import logging
+from olgram.migrations.custom import migrate
+
+logging.basicConfig(level=logging.INFO)
+
+
+if __name__ == '__main__':
+    asyncio.get_event_loop().run_until_complete(migrate())

olgram/commands/bot_actions.py

@@ -12,7 +12,7 @@ async def delete_bot(bot: Bot, call: types.CallbackQuery):
     Пользователь решил удалить бота
     """
     try:
-        await unregister_token(bot.token)
+        await unregister_token(bot.decrypted_token())
     except Unauthorized:
         # Вероятно пользователь сбросил токен или удалил бот, это уже не наши проблемы
         pass

olgram/commands/bots.py

@@ -99,7 +99,7 @@ async def bot_added(message: types.Message, state: FSMContext):
         return await on_unknown_error()
 
     user, _ = await User.get_or_create(telegram_id=message.from_user.id)
-    bot = Bot(token=token, owner=user, name=test_bot_info.username, super_chat_id=message.from_user.id)
+    bot = Bot(token=Bot.encrypted_token(token), owner=user, name=test_bot_info.username, super_chat_id=message.from_user.id)
     try:
         await bot.save()
     except IntegrityError:

olgram/migrations/custom.py

@@ -0,0 +1,35 @@
+"""Наши собственные миграции, которые нельзя описать на языке SQL и с которыми не справится TortoiseORM/Aerich"""
+
+from tortoise import transactions, Tortoise
+from olgram.settings import TORTOISE_ORM
+from olgram.models.models import MetaInfo, Bot
+import logging
+
+
+async def upgrade_1():
+    """Шифруем токены"""
+    meta_info = await MetaInfo.first()
+    if meta_info.version != 0:
+        logging.info("skip")
+        return
+
+    async with transactions.in_transaction():
+        bots = await Bot.all()
+        for bot in bots:
+            bot.token = bot.encrypted_token(bot.token)
+            await bot.save()
+        meta_info.version = 1
+        await meta_info.save()
+    logging.info("done")
+
+# Не забудь добавить миграцию в этот лист!
+_migrations = [upgrade_1]
+
+
+async def migrate():
+    logging.info("Run custom migrations...")
+    await Tortoise.init(config=TORTOISE_ORM)
+
+    for migration in _migrations:
+        logging.info(f"Migration {migration.__name__}...")
+        await migration()

olgram/migrations/models/5_20210926185420_update.sql

@@ -0,0 +1,10 @@
+-- upgrade --
+ALTER TABLE "bot" ALTER COLUMN "token" TYPE VARCHAR(200) USING "token"::VARCHAR(200);
+CREATE TABLE IF NOT EXISTS "_custom_meta_info" (
+    "id" SERIAL NOT NULL PRIMARY KEY,
+    "version" INT NOT NULL  DEFAULT 0
+);;
+INSERT INTO _custom_meta_info (id, version) VALUES (0, 0) ON CONFLICT DO NOTHING;
+-- downgrade --
+ALTER TABLE "bot" ALTER COLUMN "token" TYPE VARCHAR(50) USING "token"::VARCHAR(50);
+DROP TABLE IF EXISTS "_custom_meta_info";

olgram/models/models.py

@@ -2,11 +2,27 @@ from tortoise.models import Model
 from tortoise import fields
 from uuid import uuid4
 from textwrap import dedent
+from olgram.settings import DatabaseSettings
+
+
+class MetaInfo(Model):
+    id = fields.IntField(pk=True)
+    version = fields.IntField(default=0)
+
+    def __init__(self, **kwargs):
+        # Кажется это единственный способ сделать single-instance модель в TortoiseORM :(
+        if "id" in kwargs:
+            kwargs["id"] = 0
+        self.id = 0
+        super(MetaInfo, self).__init__(**kwargs)
+
+    class Meta:
+        table = '_custom_meta_info'
 
 
 class Bot(Model):
     id = fields.IntField(pk=True)
-    token = fields.CharField(max_length=50, unique=True)
+    token = fields.CharField(max_length=200, unique=True)
     owner = fields.ForeignKeyField("models.User", related_name="bots")
     name = fields.CharField(max_length=33)
     code = fields.UUIDField(default=uuid4, index=True)
@@ -22,6 +38,15 @@ class Bot(Model):
                                         on_delete=fields.relational.CASCADE,
                                         null=True)
 
+    def decrypted_token(self):
+        cryptor = DatabaseSettings.cryptor()
+        return cryptor.decrypt(self.token)
+
+    @classmethod
+    def encrypted_token(cls, token: str):
+        cryptor = DatabaseSettings.cryptor()
+        return cryptor.encrypt(token)
+
     async def super_chat_id(self):
         group_chat = await self.group_chat
         if group_chat:

olgram/settings.py

@@ -1,18 +1,22 @@
 from dotenv import load_dotenv
 from abc import ABC
 import os
+from olgram.utils.crypto import Cryptor
+from functools import lru_cache
 
 
 load_dotenv()
 
 
+# TODO: рефакторинг, использовать какой-нибудь lazy-config вместо своих костылей
+
 class AbstractSettings(ABC):
     @classmethod
     def _get_env(cls, parameter: str, allow_none: bool = False) -> str:
-        parameter = os.getenv(parameter, None)
-        if not parameter and not allow_none:
+        parameter_v = os.getenv(parameter, None)
+        if not parameter_v and not allow_none:
             raise ValueError(f"{parameter} not defined in ENV")
-        return parameter
+        return parameter_v
 
 
 class OlgramSettings(AbstractSettings):
@@ -99,6 +103,12 @@ class DatabaseSettings(AbstractSettings):
     def host(cls) -> str:
         return cls._get_env("POSTGRES_HOST")
 
+    @classmethod
+    @lru_cache
+    def cryptor(cls) -> Cryptor:
+        password = cls._get_env("TOKEN_ENCRYPTION_KEY")
+        return Cryptor(password)
+
 
 TORTOISE_ORM = {
     "connections": {"default": f'postgres://{DatabaseSettings.user()}:{DatabaseSettings.password()}'

olgram/utils/crypto.py

@@ -0,0 +1,16 @@
+import base64
+from Crypto.Cipher import AES
+
+
+class Cryptor:
+    def __init__(self, password: str):
+        password = password.rjust(32)[:32]
+        self._cipher = AES.new(password.encode("utf-8"), AES.MODE_ECB)
+
+    def encrypt(self, data: str) -> str:
+        if data.startswith(" "):
+            raise ValueError("Data should not start with space!")
+        return base64.b64encode(self._cipher.encrypt(data.encode("utf-8").rjust(64))).decode("utf-8")
+
+    def decrypt(self, data: str) -> str:
+        return self._cipher.decrypt(base64.b64decode(data.encode("utf-8"))).decode("utf-8").lstrip()

requirements.txt

@@ -5,3 +5,4 @@ python-dotenv~=0.17.1
 aioredis==1.3.1
 aiocache
 aiohttp
+pycrypto

server/custom.py

@@ -114,7 +114,7 @@ class CustomRequestHandler(WebhookRequestHandler):
         if not bot:
             return None
         db_bot_instance.set(bot)
-        dp = Dispatcher(AioBot(bot.token))
+        dp = Dispatcher(AioBot(bot.decrypted_token()))
 
         dp.register_message_handler(message_handler, content_types=[types.ContentType.TEXT,
                                                                     types.ContentType.CONTACT,

server/server.py

@@ -26,9 +26,9 @@ async def register_token(bot: Bot) -> bool:
     :param bot: Бот
     :return: получилось ли
     """
-    await unregister_token(bot.token)
+    await unregister_token(bot.decrypted_token())
 
-    a_bot = AioBot(bot.token)
+    a_bot = AioBot(bot.decrypted_token())
     certificate = None
     if ServerSettings.use_custom_cert():
         certificate = open(ServerSettings.public_path(), 'rb')