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

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()