1 files changed, 30 insertions, 4 deletions

diff --git a/bot/bot.py b/bot/bot.py
index 90b8141..1370c54 100644
--- a/bot/bot.py
+++ b/bot/bot.py
@@ -2,7 +2,7 @@ import io
 import os
 import logging
 import traceback
-from PIL import Image
+from PIL import Image, ExifTags
 import boto3
 from telegram import Update
 from telegram.constants import ParseMode
@@ -11,9 +11,19 @@ from peewee import *
 from uuid import uuid4
 from datetime import datetime
 import json
+import sys
 import httpx
 
-db = SqliteDatabase("database/photos.db")
+
+if os.getenv("DB_DRIVER") == "psql":
+    db = PostgresqlDatabase(os.getenv("DB_NAME"), user=os.getenv("DB_USER"), password=os.getenv("DB_PASS"),
+                            host=os.getenv("DB_HOST"), port=os.getenv("DB_PORT"))
+elif os.getenv("DB_DRIVER") == "sqlite3":
+    db = SqliteDatabase("database/photos.db")
+else:
+    print("No DB_DRIVER specified")
+    sys.exit(1)
+
 db.connect(reuse_if_open=True)
 
 
@@ -32,6 +42,7 @@ class Photo(BaseModel):
     path = CharField(max_length=256)
     caption = CharField(max_length=256)
     alt = CharField(max_length=256)
+    model = CharField(max_length=256, default="")
     createdAt = DateTimeField()
     uploadedAt = DateTimeField()
 
@@ -72,6 +83,7 @@ def write_json() -> bool:
                 "path": photo.path,
                 "caption": photo.caption,
                 "alt": photo.alt,
+                "model": photo.model,
                 "createdAt": photo.createdAt.strftime("%Y-%m-%dT%H:%M:%S.000Z"),
                 "uploadedAt": photo.uploadedAt.strftime("%Y-%m-%dT%H:%M:%S.000Z")
             })
@@ -126,6 +138,7 @@ async def process(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     else:
         return
 
+    caption = update.effective_message.caption
     await context.bot.send_message(chat_id, "Processing `{}`".format(filename), parse_mode=ParseMode.MARKDOWN_V2)
 
     img = io.BytesIO()
@@ -133,6 +146,18 @@ async def process(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
 
     try:
         im = Image.open(img)
+        exif = im.getexif()
+
+        camera_model = ""
+        shot_datetime = ""
+        if exif:
+            for key, val in exif.items():
+                if key in ExifTags.TAGS:
+                    if ExifTags.TAGS[key] == "Model":
+                        camera_model = val
+                    elif ExifTags.TAGS[key] == "DateTime":
+                        shot_datetime = val
+
         output = io.BytesIO()
         im.save(output, format="webp", lossless=False, quality=80)
 
@@ -147,9 +172,10 @@ async def process(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
                                  ratio=im.width / im.height,
                                  orientation="landscape" if im.width > im.height else "portrait",
                                  path="https://pixelstatic.jinwei.me/{}".format(key_name),
-                                 caption="",
+                                 caption=caption,
                                  alt="",
-                                 createdAt=datetime.now(),
+                                 model=camera_model,
+                                 createdAt=shot_datetime if len(shot_datetime) > 0 else datetime.now(),
                                  uploadedAt=datetime.now())
 
             output.seek(0)