bot: test peewee sqlite3

9 files changed, 139 insertions, 85 deletions

diff --git a/.gitignore b/.gitignore
index f52378c..6b3bccc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,6 +2,7 @@ foursquare/*_example.json
 config.ini
 .idea/
 fsq_poi.db
+checkinbot.db
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
diff --git a/README b/README
index b8176ad..2150cc4 100644
--- a/README
+++ b/README
@@ -11,3 +11,4 @@ TODO:
 - OAuth login to Mastodon/Pleroma
 - Delayed checkins
 - user usage count, add Telegram payment option
+- connect one Telegram user to multiple Fediverse accounts
diff --git a/bot.py b/bot.py
index 79fba5d..a059c9c 100644
--- a/bot.py
+++ b/bot.py
@@ -61,6 +61,7 @@ from config import (
 )
 from mastodon import Mastodon
 
+from dbstore.peewee_store import db, User
 
 # Enable logging
 logging.basicConfig(
@@ -72,39 +73,35 @@ logger = logging.getLogger(__name__)
 @dataclass
 class FediLoginCallbackUpdate:
     code: str
-    state: int
+    state: str
 
 
 class FediLoginCallbackContext(CallbackContext[ExtBot, dict, dict, dict]):
-    @classmethod
-    def from_update(
-            cls,
-            update: object,
-            application: "Application",
-    ) -> "FediLoginCallbackContext":
-        if isinstance(update, FediLoginCallbackUpdate):
-            return cls(application=application, user_id=update.state)
-        return super().from_update(update, application)
+    pass
 
 
 async def process_oauth_login_callback(update: FediLoginCallbackUpdate, context: FediLoginCallbackContext) -> None:
-    # query client_id and client_secret from in memory database
-    client_id = ""
-    client_secret = ""
-    home_instance = ""
-
-    mastodon_client = Mastodon(client_id=client_id, client_secret=client_secret, api_base_url=home_instance)
-    access_token = mastodon_client.log_in(
-        code=update.code,
-        redirect_uri="{}{}".format(BOT_DOMAIN, FEDI_LOGIN_CALLBACK_URL),
-        scopes=BOT_SCOPE
-    )
+    state = update.state
+
+    with db.connection_context():
+        user = User.get(User.state == state)
 
-    # TODO
-    # save access_token to database
+        client_id = user.client_id
+        client_secret = user.client_secret
+        home_instance = user.home_instance
+
+        if len(user.access_key) == 0:
+            mastodon_client = Mastodon(client_id=client_id, client_secret=client_secret, api_base_url=home_instance)
+            access_token = mastodon_client.log_in(
+                code=update.code,
+                redirect_uri="{}{}".format(BOT_DOMAIN, FEDI_LOGIN_CALLBACK_URL),
+                scopes=BOT_SCOPE
+            )
+            user.access_key = access_token
+            user.save()
 
-    text = "You have successfully logged in to your Mastodon account!"
-    await context.bot.send_message(chat_id=update.state, text=text)
+        text = "You have successfully logged in to your Mastodon account!"
+        await context.bot.send_message(chat_id=user.telegram_user_id, text=text)
 
 
 async def main() -> None:
@@ -112,7 +109,7 @@ async def main() -> None:
     # Here we set updater to None because we want our custom webhook server to handle the updates
     # and hence we don't need an Updater instance
     application = (
-        Application.builder().token(BOT_TOKEN).context_types(context_types).build()
+        Application.builder().updater(None).token(BOT_TOKEN).context_types(context_types).build()
     )
 
     checkin_handler = ConversationHandler(
@@ -171,7 +168,7 @@ async def main() -> None:
         """
         try:
             code = request.query_params["code"]
-            state = int(request.query_params.get("state"))
+            state = request.query_params.get("state")
         except KeyError:
             return PlainTextResponse(
                 status_code=HTTPStatus.BAD_REQUEST,
diff --git a/callback.py b/callback.py
index c7cb60c..4fe0540 100644
--- a/callback.py
+++ b/callback.py
@@ -7,13 +7,23 @@ from telegram.error import BadRequest
 from telegram.ext import CallbackContext
 
 from command import *
-from dbstore.dbm_store import get_loc
+from dbstore.peewee_store import get_poi_by_fsq_id
 from foursquare.poi import OSM_ENDPOINT
 from foursquare.poi import query_poi
-# from toot import mastodon_client
 from config import BOT_SCOPE
+from dbstore.peewee_store import User, db
+import uuid
 
-mastodon_client = None
+
+def generate_uuid():
+    return str(uuid.uuid4())
+
+
+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)
 
 
 def generate_toot_text(poi_name, poi_locality, poi_region, poi_lat, poi_lon):
@@ -41,6 +51,7 @@ async def get_img_file_bytes(telegram_media_file):
 async def process_media_group(context: CallbackContext):
     context.job.data = cast(List[MsgDict], context.job.data)
 
+    mastodon_client = get_mastodon_client(context.user_data["user_id"])
     media_id = []
     chat_id = context.job.data[0].get("chat_id")
     for media_dict in context.job.data:
@@ -61,9 +72,6 @@ async def process_media_group(context: CallbackContext):
 
 
 async def callback_generate_fedi_login_url(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
-    # generate fedi OAuth login url
-    global mastodon_client
-
     home_instance = update.effective_message.text
     client_id, client_secret = Mastodon.create_app(
         "Checkin.bot",
@@ -71,13 +79,25 @@ async def callback_generate_fedi_login_url(update: Update, context: ContextTypes
         redirect_uris="{}{}".format(BOT_DOMAIN, FEDI_LOGIN_CALLBACK_URL),
         api_base_url=home_instance,
     )
+    print("client_id: {}".format(client_id))
+    print("client_secret: {}".format(client_secret))
+
     m = Mastodon(client_id=client_id, client_secret=client_secret, api_base_url=home_instance)
 
-    # TODO
-    # generate random string as OAuth state
+    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()
+
     oauth_url = m.auth_request_url(redirect_uris="{}{}".format(BOT_DOMAIN, FEDI_LOGIN_CALLBACK_URL),
                                                  scopes=BOT_SCOPE,
-                                                 state=update.effective_user.id)
+                                                 state=state)
 
     await update.message.reply_text(PROMPT_FEDI_LOGIN,
                                     reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Login", url=oauth_url)]]),
@@ -93,9 +113,9 @@ async def callback_location_sharing(update: Update, context: ContextTypes.DEFAUL
         context.user_data["latitude"] = update.message.venue.location.latitude
         context.user_data["longitude"] = update.message.venue.location.longitude
 
-        poi = get_loc(context.user_data.get("fsq_id"))
+        poi = get_poi_by_fsq_id(context.user_data.get("fsq_id"))
         content = generate_toot_text(poi["name"], poi["locality"], poi["region"], poi["latitude"], poi["longitude"])
-        status = mastodon_client.status_post(content, visibility=DEFAULT_TOOT_VISIBILITY, media_ids=[])
+        status = get_mastodon_client(update.effective_user.id).status_post(content, visibility=DEFAULT_TOOT_VISIBILITY, media_ids=[])
 
         context.user_data[KEY_TOOT_STATUS_ID] = status["id"]
         context.user_data[KEY_TOOT_STATUS_CONTENT] = content
@@ -165,12 +185,12 @@ async def callback_skip_location_keyword_search(update: Update, context: Context
 async def _process_location_selection(context: ContextTypes.DEFAULT_TYPE) -> int:
     poi_name = context.user_data.get("poi_name")
     if context.user_data.get("fsq_id") is not None:
-        poi = get_loc(context.user_data.get("fsq_id"))
+        poi = get_poi_by_fsq_id(context.user_data.get("fsq_id"))
         content = generate_toot_text(poi["name"], poi["locality"], poi["region"], poi["latitude"], poi["longitude"])
     else:
         content = generate_toot_text(poi_name, "", "", context.user_data.get("latitude"), context.user_data.get("longitude"))
 
-    status = mastodon_client.status_post(content, visibility=DEFAULT_TOOT_VISIBILITY, media_ids=[])
+    status = get_mastodon_client(context.user_data["user_id"]).status_post(content, visibility=DEFAULT_TOOT_VISIBILITY, media_ids=[])
 
     context.user_data[KEY_TOOT_STATUS_ID] = status["id"]
     context.user_data[KEY_TOOT_STATUS_CONTENT] = content
@@ -192,6 +212,7 @@ async def callback_location_confirmation(update: Update, context: ContextTypes.D
     query = update.callback_query
     await query.answer()
     context.user_data["fsq_id"] = query.data
+    context.user_data["user_id"] = update.effective_user.id
     await query.delete_message()
 
     context.user_data["chat_id"] = update.effective_chat.id
@@ -202,6 +223,7 @@ async def callback_location_confirmation(update: Update, context: ContextTypes.D
 async def callback_manual_location(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
     context.user_data["poi_name"] = update.effective_message.text
     context.user_data["chat_id"] = update.effective_chat.id
+    context.user_data["user_id"] = update.effective_user.id
 
     return await _process_location_selection(context)
 
@@ -219,7 +241,7 @@ async def callback_add_comment(update: Update, context: ContextTypes.DEFAULT_TYP
     await context.bot.delete_message(update.effective_chat.id, context.user_data.get(PROMPT_ADD_COMMENT))
 
     comment = update.effective_message.text
-    mastodon_client.status_update(id=context.user_data.get(KEY_TOOT_STATUS_ID),
+    get_mastodon_client(update.effective_user.id).status_update(id=context.user_data.get(KEY_TOOT_STATUS_ID),
                                   status=f"{comment} " + context.user_data.get(KEY_TOOT_STATUS_CONTENT))
     context.user_data[KEY_TOOT_STATUS_CONTENT] = f"{comment} " + context.user_data.get(KEY_TOOT_STATUS_CONTENT)
 
@@ -243,6 +265,7 @@ async def callback_add_media(update: Update, context: CallbackContext):
         if "not found" in str(e.message):
             pass
 
+    mastodon_client = get_mastodon_client(update.effective_user.id)
     status_id = context.user_data.get(KEY_TOOT_STATUS_ID)
     status_content = context.user_data.get(KEY_TOOT_STATUS_CONTENT)
 
diff --git a/command.py b/command.py
index 837d3cd..e45a2e2 100644
--- a/command.py
+++ b/command.py
@@ -15,13 +15,6 @@ async def start_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> i
 
 
 async def fedi_login_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
-    # generate fedi OAuth login url
-
-    # mastodon_client = Mastodon(client_id=MASTODON_CLIENT_ID_FILE, api_base_url=TOOT_API_BASE_URL)
-    # oauth_url = mastodon_client.auth_request_url(redirect_uris="{}/{}".format(BOT_DOMAIN, FEDI_LOGIN_CALLBACK_URL),
-    #                                              scopes=['write:media', 'write:statuses'])
-    #
-    # await update.message.reply_text(PROMPT_FEDI_LOGIN.format(oauth_url), parse_mode=ParseMode.MARKDOWN)
     await update.message.reply_text(PROMPT_FEDI_LOGIN_WHERE_IS_INSTANCE, parse_mode=ParseMode.MARKDOWN)
     return FEDI_LOGIN
 
diff --git a/dbstore/dbm_store.py b/dbstore/dbm_store.py
deleted file mode 100644
index fedb505..0000000
--- a/dbstore/dbm_store.py
+++ /dev/null
@@ -1,34 +0,0 @@
-import dbm
-
-db = None
-store_file = "fsq_poi.db"
-
-
-def get_loc(fsq_id):
-    global db
-    if db is None:
-        db = dbm.open(store_file, 'c')
-    if fsq_id in db:
-        res = db[fsq_id].decode("utf-8").split("|")
-        return {
-            "name": res[0],
-            "locality": res[1],
-            "region": res[2],
-            "latitude": res[3],
-            "longitude": res[4],
-            "osm_url": res[5],
-        }
-    else:
-        return None
-
-
-def store_loc(loc):
-    global db
-    if db is None:
-        db = dbm.open(store_file, 'c')
-    db[loc["fsq_id"]] = "{}|{}|{}|{}|{}|{}".format(loc["name"],
-                                                   loc["locality"],
-                                                   loc["region"],
-                                                   loc["latitude"],
-                                                   loc["longitude"],
-                                                   loc["osm_url"])
diff --git a/dbstore/peewee_store.py b/dbstore/peewee_store.py
new file mode 100644
index 0000000..37b8c01
--- /dev/null
+++ b/dbstore/peewee_store.py
@@ -0,0 +1,59 @@
+from peewee import *
+
+db = SqliteDatabase("checkinbot.db")
+db.connect(reuse_if_open=True)
+
+
+class BaseModel(Model):
+    class Meta:
+        database = db
+
+
+class User(BaseModel):
+    telegram_user_id = CharField(unique=True, primary_key=True)
+    access_key = CharField(max_length=64)
+    home_instance = CharField(max_length=128)
+    state = CharField(max_length=128)
+    client_id = CharField(max_length=128)
+    client_secret = CharField(max_length=128)
+
+
+class Location(BaseModel):
+    fsq_id = CharField(unique=True, primary_key=True)
+    name = CharField(max_length=128)
+    locality = CharField(max_length=128)
+    region = CharField(max_length=128)
+    latitude = CharField(max_length=128)
+    longitude = CharField(max_length=128)
+
+
+with db.connection_context():
+    db.create_tables([User, Location])
+
+
+def get_poi_by_fsq_id(fsq_id) -> dict:
+    with db.connection_context():
+        try:
+            poi = Location.get(Location.fsq_id == fsq_id)
+            return {
+                "name": poi.name,
+                "locality": poi.locality,
+                "region": poi.region,
+                "latitude": poi.latitude,
+                "longitude": poi.longitude,
+            }
+        except DoesNotExist:
+            return {}
+
+
+def create_poi(poi: dict):
+    with db.connection_context():
+        poi = Location.create(
+            fsq_id=poi["fsq_id"],
+            name=poi["name"],
+            locality=poi["locality"],
+            region=poi["region"],
+            latitude=poi["latitude"],
+            longitude=poi["longitude"],
+        )
+        poi.save()
diff --git a/foursquare/poi.py b/foursquare/poi.py
index 5e67d1c..45b4fa9 100644
--- a/foursquare/poi.py
+++ b/foursquare/poi.py
@@ -3,7 +3,7 @@ import json
 import requests
 
 from config import FSQ_API_KEY
-from dbstore.dbm_store import store_loc
+from dbstore.peewee_store import create_poi
 
 POI_API_ENDPOINT = "https://api.foursquare.com/v3/places/nearby?ll={}%2C{}&limit=10"
 POI_SEARCH_API_ENDPOINT = "https://api.foursquare.com/v3/places/search?query={}&ll={}%2C{}&radius=2000&limit=10"
@@ -33,8 +33,7 @@ def query_poi(search, latitude, longitude):
             "osm_url": OSM_ENDPOINT.format(poi["geocodes"]["main"]["latitude"], poi["geocodes"]["main"]["longitude"])
         }
         locations.append(loc)
-        print(loc)
-        store_loc(loc)
+        create_poi(loc)
 
     return locations
 
diff --git a/requirements.txt b/requirements.txt
index c031383..c3166b6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,8 +1,12 @@
 anyio==3.6.2
+APScheduler==3.10.0
+asgiref==3.6.0
 blurhash==1.1.4
 certifi==2022.12.7
 charset-normalizer==3.0.1
+click==8.1.3
 decorator==5.1.1
+greenlet==2.0.2
 h11==0.14.0
 h2==4.1.0
 hpack==4.0.0
@@ -11,11 +15,22 @@ httpx==0.23.3
 hyperframe==6.0.1
 idna==3.4
 Mastodon.py==1.8.0
+peewee==3.15.4
+Pillow==9.4.0
 python-dateutil==2.8.2
 python-magic==0.4.27
 python-telegram-bot==20.1
+pytz==2022.7.1
+pytz-deprecation-shim==0.1.0.post0
 requests==2.28.2
 rfc3986==1.5.0
 six==1.16.0
 sniffio==1.3.0
+SQLAlchemy==2.0.4
+sqlparse==0.4.3
+starlette==0.25.0
+typing_extensions==4.5.0
+tzdata==2022.7
+tzlocal==4.2
 urllib3==1.26.14
+uvicorn==0.20.0