Merge branch 'feature/bot'

add python telegram bot to upload photos to cloudflare r2

3 files changed, 200 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
index fafa5b5..f157e65 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,3 +3,4 @@ _site/
 .DS_Store
 .ruby-version
 .jekyll-cache/
+.venv/
diff --git a/bot.py b/bot.py
new file mode 100644
index 0000000..20529c0
--- /dev/null
+++ b/bot.py
@@ -0,0 +1,178 @@
+import io
+import os
+import logging
+import traceback
+from PIL import Image
+import boto3
+from telegram import Update
+from telegram.constants import ParseMode
+from telegram.ext import Application, CommandHandler, ContextTypes, MessageHandler, filters
+from peewee import *
+from uuid import uuid4
+from datetime import datetime
+import json
+
+db = SqliteDatabase("database/photos.db")
+db.connect(reuse_if_open=True)
+
+
+class BaseModel(Model):
+    class Meta:
+        database = db
+
+
+class Photo(BaseModel):
+    guid = CharField(unique=True, primary_key=True)
+    fileId = CharField(max_length=256)
+    width = IntegerField()
+    height = IntegerField()
+    ratio = FloatField()
+    orientation = CharField(max_length=128)
+    path = CharField(max_length=256)
+    caption = CharField(max_length=256)
+    alt = CharField(max_length=256)
+    createdAt = DateTimeField()
+    uploadedAt = DateTimeField()
+
+
+with db.connection_context():
+    db.create_tables([Photo])
+
+logging.basicConfig(
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+
+logger = logging.getLogger(__name__)
+
+
+async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+    await update.message.reply_text("This is a bot to output image in square shape")
+
+
+bucket_name = "pixel-jinwei-me"
+cf_account_id = os.getenv("CF_ACCOUNT_ID")
+aws_access_key_id = os.getenv("CF_R2_KEY_ID")
+aws_secret_access_key = os.getenv("CF_R2_ACCESS_KEY_SECRET")
+
+
+def write_json() -> bool:
+    with db.connection_context():
+        photos = Photo.select()
+        results = []
+        for photo in photos:
+            results.append({
+                "guid": photo.guid,
+                "fileId": photo.fileId,
+                "width": photo.width,
+                "height": photo.height,
+                "ratio": photo.ratio,
+                "orientation": photo.orientation,
+                "path": photo.path,
+                "caption": photo.caption,
+                "alt": photo.alt,
+                "createdAt": photo.createdAt.strftime("%Y-%m-%dT%H:%M:%S.000Z"),
+                "uploadedAt": photo.uploadedAt.strftime("%Y-%m-%dT%H:%M:%S.000Z")
+            })
+        with open("database/photos.json", "w") as f:
+            f.write(json.dumps(results))
+    return True
+
+
+def upload_to_s3(key_name: str, file: bytes):
+    endpoint_ca_central = "https://{}.r2.cloudflarestorage.com".format(cf_account_id)
+
+    client = boto3.client("s3",
+                          region_name="auto",
+                          endpoint_url=endpoint_ca_central,
+                          aws_access_key_id=aws_access_key_id,
+                          aws_secret_access_key=aws_secret_access_key)
+
+    response = client.list_buckets()
+    buckets = [b["Name"] for b in response['Buckets']]
+
+    if bucket_name not in buckets:
+        print("{} doesn't exist".format(bucket_name))
+        print('Existing buckets:')
+        for bucket in response['Buckets']:
+            print(f' {bucket["Name"]}')
+        return False
+
+    response = client.put_object(Body=file, Bucket=bucket_name, Key=key_name, ContentType="image/webp")
+    status = response["ResponseMetadata"]["HTTPStatusCode"]
+    if status == 200:
+        print("upload {} to {} succeed".format(key_name, bucket_name))
+        return True
+    else:
+        print("upload {} to {} failed, status: {}".format(key_name, bucket_name, status))
+        return False
+
+
+async def process(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+    chat_id = update.message.chat_id
+
+    if update.message.document is not None:
+        names = update.message.document.file_name.split(".")
+        file_ext = names[1]
+        filename = names[0]
+
+        if str.upper(file_ext) not in ("JPG", "JPEG", "PNG"):
+            await context.bot.send_message(chat_id, "Image extension `{}` not supported".format(file_ext),
+                                           parse_mode=ParseMode.MARKDOWN_V2)
+            return
+
+        file = await update.message.effective_attachment.get_file()
+    else:
+        return
+
+    await context.bot.send_message(chat_id, "Processing `{}`".format(filename), parse_mode=ParseMode.MARKDOWN_V2)
+
+    img = io.BytesIO()
+    await file.download_to_memory(img)
+
+    try:
+        im = Image.open(img)
+        output = io.BytesIO()
+        im.save(output, format="webp", lossless=True, quality=100)
+
+        now = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+        key_name = "{}-{}.webp".format(now, filename)
+
+        with db.connection_context():
+            photo = Photo.create(guid=str.upper(str(uuid4())),
+                                 fileId=key_name,
+                                 width=im.width,
+                                 height=im.height,
+                                 ratio=im.width / im.height,
+                                 orientation="landscape" if im.width > im.height else "portrait",
+                                 path="https://pixelstatic.jinwei.me/{}".format(key_name),
+                                 caption="dsadas",
+                                 alt="dsada",
+                                 createdAt=datetime.now(),
+                                 uploadedAt=datetime.now())
+
+            output.seek(0)
+            upload_to_s3(photo.fileId, output.read())
+
+        write_json()
+        await update.message.reply_markdown_v2(text="Sending processed result")
+
+        await context.bot.send_document(chat_id=update.message.chat_id,
+                                        filename="{}-result.{}".format(filename, file_ext),
+                                        document=output.getvalue())
+
+    except Exception as e:
+        await update.message.reply_markdown_v2(text="Error:\n```{}```".format(traceback.format_exc()))
+
+
+def main() -> None:
+    tg_token = os.getenv("TG_TOKEN")
+    application = Application.builder().token(tg_token).build()
+
+    application.add_handler(CommandHandler("start", start))
+    application.add_handler(MessageHandler(filters.ATTACHMENT & ~filters.COMMAND, process))
+
+    application.run_polling()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..788075c
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,21 @@
+anyio==3.6.2
+boto3==1.26.81
+botocore==1.29.81
+certifi==2022.12.7
+h11==0.14.0
+h2==4.1.0
+hpack==4.0.0
+httpcore==0.16.3
+httpx==0.23.3
+hyperframe==6.0.1
+idna==3.4
+jmespath==1.0.1
+peewee==3.16.0
+Pillow==9.4.0
+python-dateutil==2.8.2
+python-telegram-bot==20.1
+rfc3986==1.5.0
+s3transfer==0.6.0
+six==1.16.0
+sniffio==1.3.0
+urllib3==1.26.14