From e1e8c54eafd61d31aa520593ece8a68d79dbdeeb Mon Sep 17 00:00:00 2001 From: clarkzjw Date: Tue, 28 Feb 2023 23:27:35 -0800 Subject: add python telegram bot to upload photos to cloudflare r2 --- .gitignore | 1 + bot.py | 178 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 21 +++++++ 3 files changed, 200 insertions(+) create mode 100644 bot.py create mode 100644 requirements.txt 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 -- cgit v1.2.3