add caption and exif

5 files changed, 57 insertions, 10 deletions

diff --git a/.gitignore b/.gitignore
index fbc84df..ab3415e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,5 @@
+.env
+pg.yml
 .idea/
 bot/database/
 env
diff --git a/bot/.dockerignore b/bot/.dockerignore
new file mode 100644
index 0000000..a2a8dea
--- /dev/null
+++ b/bot/.dockerignore
@@ -0,0 +1 @@
+.env/
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)
diff --git a/bot/docker-compose.yaml b/bot/docker-compose.yaml
index eba14ea..4d22739 100644
--- a/bot/docker-compose.yaml
+++ b/bot/docker-compose.yaml
@@ -1,14 +1,10 @@
 version: '3'
 services:
-  photobot:
+  bot:
     build:
         context: .
         dockerfile: Dockerfile
     restart: always
+    network_mode: "host"
     env_file:
       - .env
-    volumes:
-      - sqlite3:/usr/src/app/database/
-
-volumes:
-    sqlite3:
diff --git a/bot/migrations/0001-add_exif.py b/bot/migrations/0001-add_exif.py
new file mode 100644
index 0000000..0eff797
--- /dev/null
+++ b/bot/migrations/0001-add_exif.py
@@ -0,0 +1,22 @@
+import os
+import sys
+from playhouse.migrate import *
+
+
+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"))
+    migrator = PostgresqlMigrator(db)
+elif os.getenv("DB_DRIVER") == "sqlite3":
+    db = SqliteDatabase("database/photos.db")
+    migrator = SqliteMigrator(db)
+else:
+    print("No DB_DRIVER specified")
+    sys.exit(1)
+
+model = CharField(max_length=256, default="")
+
+# Run the migration, specifying the database table, field name and field.
+migrate(
+    migrator.add_column('photo', 'model', model),
+)