bot: encrypt access_key with cryptography.fernet library

7 files changed, 36 insertions, 19 deletions

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