From 2e61640f5a10c438463ff29b89d9e8dfa79a2ee5 Mon Sep 17 00:00:00 2001 From: mihalin Date: Sun, 26 Sep 2021 20:36:05 +0300 Subject: [PATCH] =?UTF-8?q?=D0=A8=D0=B8=D1=84=D1=80=D0=BE=D0=B2=D0=B0?= =?UTF-8?q?=D0=BD=D0=B8=D0=B5=20=D1=82=D0=BE=D0=BA=D0=B5=D0=BD=D0=BE=D0=B2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 3 +- docker-entrypoint.sh | 1 + example.env | 2 ++ main.py | 2 +- migrate.py | 9 +++++ olgram/commands/bot_actions.py | 2 +- olgram/commands/bots.py | 2 +- olgram/migrations/custom.py | 35 +++++++++++++++++++ .../models/5_20210926185420_update.sql | 10 ++++++ olgram/models/models.py | 27 +++++++++++++- olgram/settings.py | 16 +++++++-- olgram/utils/crypto.py | 16 +++++++++ requirements.txt | 3 +- server/custom.py | 2 +- server/server.py | 4 +-- 15 files changed, 122 insertions(+), 12 deletions(-) create mode 100644 migrate.py create mode 100644 olgram/migrations/custom.py create mode 100644 olgram/migrations/models/5_20210926185420_update.sql create mode 100644 olgram/utils/crypto.py diff --git a/.gitignore b/.gitignore index fa373c5..49a2a2e 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,5 @@ venv __pycache__ *.pyc config.json -docker-compose-release.yaml \ No newline at end of file +docker-compose-release.yaml +docs/build \ No newline at end of file diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh index efb3d88..b5561da 100755 --- a/docker-entrypoint.sh +++ b/docker-entrypoint.sh @@ -11,4 +11,5 @@ fi sleep 10 aerich upgrade +python migrate.py python main.py diff --git a/example.env b/example.env index d2435b8..9f8c2d3 100644 --- a/example.env +++ b/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 diff --git a/main.py b/main.py index 311a625..2a4a254 100644 --- a/main.py +++ b/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() diff --git a/migrate.py b/migrate.py new file mode 100644 index 0000000..6214856 --- /dev/null +++ b/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()) diff --git a/olgram/commands/bot_actions.py b/olgram/commands/bot_actions.py index 5471b1c..ac9a9e6 100644 --- a/olgram/commands/bot_actions.py +++ b/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 diff --git a/olgram/commands/bots.py b/olgram/commands/bots.py index 76a6b52..5a3f19d 100644 --- a/olgram/commands/bots.py +++ b/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: diff --git a/olgram/migrations/custom.py b/olgram/migrations/custom.py new file mode 100644 index 0000000..da1b3d4 --- /dev/null +++ b/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() diff --git a/olgram/migrations/models/5_20210926185420_update.sql b/olgram/migrations/models/5_20210926185420_update.sql new file mode 100644 index 0000000..2cf267c --- /dev/null +++ b/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"; diff --git a/olgram/models/models.py b/olgram/models/models.py index 2ac7271..aab4cbd 100644 --- a/olgram/models/models.py +++ b/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: diff --git a/olgram/settings.py b/olgram/settings.py index c5f6515..499392f 100644 --- a/olgram/settings.py +++ b/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()}' diff --git a/olgram/utils/crypto.py b/olgram/utils/crypto.py new file mode 100644 index 0000000..394965e --- /dev/null +++ b/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() diff --git a/requirements.txt b/requirements.txt index 118b835..c5784a5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ aerich==0.5.4 python-dotenv~=0.17.1 aioredis==1.3.1 aiocache -aiohttp \ No newline at end of file +aiohttp +pycrypto \ No newline at end of file diff --git a/server/custom.py b/server/custom.py index ebd6b4c..170048f 100644 --- a/server/custom.py +++ b/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, diff --git a/server/server.py b/server/server.py index 01b8b37..97bbb49 100644 --- a/server/server.py +++ b/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')