From 4cee8500a51aff6d0a445a3a6259cafed92d4845 Mon Sep 17 00:00:00 2001 From: clarkzjw Date: Thu, 23 Feb 2023 15:37:48 -0800 Subject: bot: encrypt access_key with cryptography.fernet library --- README.md | 1 - bot.py | 5 +++-- callback.py | 22 +++++++++++----------- config.ini.example | 1 + config.py | 2 +- dbstore/peewee_store.py | 6 ++---- util.py | 18 ++++++++++++++++++ 7 files changed, 36 insertions(+), 19 deletions(-) create mode 100644 util.py diff --git a/README.md b/README.md index 4d71069..4881eb8 100644 --- a/README.md +++ b/README.md @@ -24,4 +24,3 @@ Foursquare Swarm like Telegram bot to check in at places and post to Fediverse ( - [ ] add Telegram payment option to enable raw GPS checkins - [ ] i18n - [ ] Anonymized usage analysis, monitoring (Grafana/Prometheus) - diff --git a/bot.py b/bot.py index 45b2af7..7d10ddd 100644 --- a/bot.py +++ b/bot.py @@ -4,7 +4,7 @@ import asyncio import logging from dataclasses import dataclass from http import HTTPStatus -from config import BOT_TOKEN, TELEGRAM_WEBHOOK_URL, HEALTHCHECK_URL, FEDI_LOGIN_CALLBACK_URL, BOT_DOMAIN, BOT_PORT +from config import BOT_TOKEN, TELEGRAM_WEBHOOK_URL, HEALTHCHECK_URL, FEDI_LOGIN_CALLBACK_URL, BOT_DOMAIN, BOT_PORT, ENCRYPT_KEY import uvicorn from starlette.applications import Starlette from starlette.requests import Request @@ -61,6 +61,7 @@ from config import ( from prompt.string import PROMPT_CHOOSE_ACTION from mastodon import Mastodon from dbstore.peewee_store import db, User, get_user_by_state +from util import encrypt, decrypt logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO @@ -104,7 +105,7 @@ async def process_oauth_login_callback(update: FediLoginCallbackUpdate, context: redirect_uri="{}{}".format(BOT_DOMAIN, FEDI_LOGIN_CALLBACK_URL), scopes=BOT_SCOPE ) - user.access_key = access_token + user.access_key = encrypt(access_token, ENCRYPT_KEY) user.save() text = "You have successfully logged in to your Mastodon account!" diff --git a/callback.py b/callback.py index d968523..a1b4b2d 100644 --- a/callback.py +++ b/callback.py @@ -9,10 +9,11 @@ from command import * from dbstore.peewee_store import get_poi_by_fsq_id from foursquare.poi import OSM_ENDPOINT from foursquare.poi import query_poi -from config import BOT_SCOPE +from config import BOT_SCOPE, ENCRYPT_KEY from dbstore.peewee_store import User, db, TOOT_VISIBILITY_PRIVATE, TOOT_VISIBILITY_PUBLIC, TOOT_VISIBILITY_UNLISTED import uuid from mastodon import Mastodon +from util import decrypt def generate_uuid(): @@ -23,7 +24,7 @@ def get_mastodon_client(user_id: int): with db.connection_context(): user = User.get(User.telegram_user_id == user_id) if user.home_instance and user.access_key: - return Mastodon(access_token=user.access_key, api_base_url=user.home_instance) + return Mastodon(access_token=decrypt(user.access_key, ENCRYPT_KEY), api_base_url=user.home_instance) def generate_toot_text(poi_name, poi_locality, poi_region, poi_lat, poi_lon): @@ -85,21 +86,20 @@ async def callback_generate_fedi_login_url(update: Update, context: ContextTypes user_id = update.effective_user.id state = generate_uuid() - db.connect() - u = User.get_or_none(telegram_user_id=user_id) - if u is None: - u = User.create(telegram_user_id=user_id, access_key="", home_instance=home_instance, - client_id=client_id, client_secret=client_secret, state=state) - u.save() - db.close() + with db.connection_context(): + u = User.get_or_none(telegram_user_id=user_id) + if u is None: + u = User.create(telegram_user_id=user_id, access_key="", home_instance=home_instance, + client_id=client_id, client_secret=client_secret, state=state) + u.save() oauth_url = m.auth_request_url(redirect_uris="{}{}".format(BOT_DOMAIN, FEDI_LOGIN_CALLBACK_URL), scopes=BOT_SCOPE, state=state) msg = await update.message.reply_text(PROMPT_FEDI_LOGIN, - reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Login", url=oauth_url)]]), - parse_mode=ParseMode.MARKDOWN) + reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Login", url=oauth_url)]]), + parse_mode=ParseMode.MARKDOWN) context.user_data[PROMPT_FEDI_LOGIN] = msg.message_id return FEDI_LOGIN diff --git a/config.ini.example b/config.ini.example index 2efd6d5..e3a9121 100644 --- a/config.ini.example +++ b/config.ini.example @@ -1,3 +1,4 @@ [DEFAULT] BOT_TOKEN = FOURSQUARE_API_KEY = +ENCRYPT_KEY = diff --git a/config.py b/config.py index b2025ca..83ee37a 100644 --- a/config.py +++ b/config.py @@ -8,7 +8,7 @@ config.read("config.ini") BOT_TOKEN = config["DEFAULT"]["BOT_TOKEN"] FSQ_API_KEY = config["DEFAULT"]["FOURSQUARE_API_KEY"] - +ENCRYPT_KEY = config["DEFAULT"]["ENCRYPT_KEY"] MEDIA_GROUP_TIMEOUT = 3 diff --git a/dbstore/peewee_store.py b/dbstore/peewee_store.py index d124de2..445816e 100644 --- a/dbstore/peewee_store.py +++ b/dbstore/peewee_store.py @@ -16,10 +16,8 @@ class BaseModel(Model): class User(BaseModel): telegram_user_id = CharField(unique=True, primary_key=True) - # TODO: - # add AES encryption to access_key - access_key = CharField(max_length=64) - home_instance = CharField(max_length=128) + access_key = CharField(max_length=256) + home_instance = CharField(max_length=256) state = CharField(max_length=128) client_id = CharField(max_length=128) client_secret = CharField(max_length=128) diff --git a/util.py b/util.py new file mode 100644 index 0000000..61efa0b --- /dev/null +++ b/util.py @@ -0,0 +1,18 @@ +from cryptography.fernet import Fernet + + +def encrypt(input: str, key: str) -> str: + f = Fernet(key) + return f.encrypt(bytes(input, 'utf-8')).decode('utf-8') + + +def decrypt(input: str, key: str) -> str: + f = Fernet(key) + return f.decrypt(input).decode('utf-8') + + +# if __name__ == "__main__": +# key = Fernet.generate_key().decode('utf-8') +# print(key) +# print(encrypt("Hello World!", key)) +# print(decrypt(encrypt("Hello World!", key), key)) -- cgit v1.2.3