tanabata/bot/tfm-tb.py

#!../venv/bin/python3

import telebot
import sys
import os
import atexit
import signal
from time import sleep
from socket import getaddrinfo
from requests import get
from subprocess import check_output
import logging as log
from json import loads as ljson
from datetime import datetime
import pytz

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import api.tfm_api as tfm_api

# set logger
log.basicConfig(
	level=log.INFO,
	filename="/var/log/tfm/tfm-tb.log",
	filemode="a",
	format="%(asctime)s | %(threadName)s | %(levelname)s | %(message)s"
)

# actions to do on exit
exit_actions = []
defer = exit_actions.append
def finalize(*args):
	exec('\n'.join(exit_actions))
	os._exit(0)
atexit.register(finalize)
signal.signal(signal.SIGTERM, finalize)
signal.signal(signal.SIGINT, finalize)

# initialize TFM API
try:
	tfm_api.Initialize()
	tfm = tfm_api.TSession("af6dde9b-7be1-46f2-872e-9a06ce3d3358")
except Exception as e:
	log.critical(f"failed to initialize TFM API: {str(e)}")
	exit(1)

# initialize bot
tfm_tb = telebot.TeleBot(tfm_api.conf["Telebot"]["Token"])

TZ = pytz.timezone("Europe/Moscow")


# check if user is authorized and if chat exists in db
def check_chat(message):
	return message.from_user.id == 924090228


# file handler
@tfm_tb.message_handler(content_types=['document', 'photo', 'audio', 'video', 'voice', 'animation'])
def file_handler(message):
	if not check_chat(message):
		return
	notes = None
	orig_name = None
	if message.forward_from_chat:
		notes = f"Telegram origin: \"{message.forward_from_chat.title}\" ({message.forward_from_chat.username})"
	if message.photo:
		fname = f"{message.photo[-1].file_unique_id}"
		log.info(f"got photo '{fname}'")
		file_info = tfm_tb.get_file(message.photo[-1].file_id)
		file_path = os.path.join(tfm_api.conf["Telebot"]["SaveFolder"], fname)
		with open(file_path, 'wb') as f:
			f.write(tfm_tb.download_file(file_info.file_path))
	elif message.video:
		fname = f"{message.video.file_unique_id}"
		log.info(f"got video '{fname}'")
		file_info = tfm_tb.get_file(message.video.file_id)
		file_path = os.path.join(tfm_api.conf["Telebot"]["SaveFolder"], fname)
		with open(file_path, 'wb') as f:
			f.write(tfm_tb.download_file(file_info.file_path))
	else:
		file = None
		if message.document:
			file = message.document
		elif message.animation:
			file = message.animation
		else:
			tfm_tb.reply_to(message, "Unsupported file type.")
			return
		log.info(f"got file '{file.file_name}'")
		orig_name = file.file_name
		file_info = tfm_tb.get_file(file.file_id)
		file_path = os.path.join(tfm_api.conf["Telebot"]["SaveFolder"], f"{file.file_unique_id}")
	try:
		with open(file_path, 'wb') as f:
			f.write(tfm_tb.download_file(file_info.file_path))
		log.info(f"downloaded file '{file_path}'")
		exif = ljson(os.popen(f"exiftool -json \"{file_path}\"").read())[0]
		dt = exif["FileModifyDate"]
		if "SubSecCreateDate" in exif.keys():
			dt = exif["SubSecCreateDate"]
		if '.' in dt:
			dt = datetime.strptime(dt, "%Y:%m:%d %H:%M:%S.%f%z")
		else:
			dt = datetime.strptime(dt, "%Y:%m:%d %H:%M:%S%z")
		file_id, ext = tfm.add_file(file_path, dt, notes, orig_name=orig_name)
		tfm.add_file_to_tag(file_id, "d6d8129a-984d-4451-8c83-d04523ced8a8")
	except Exception as e:
		tfm_tb.reply_to(message, "Error: %s" % str(e).split('\n')[0])
		log.info(f"Error: %s" % str(e).split('\n')[0])
		os.remove(file_path)
		log.info(f"removed file '{file_path}'")
	else:
		tfm_tb.reply_to(message, "File saved.")
		log.info(f"imported file '{file_path}'")


# folder scanner
@tfm_tb.message_handler(commands=['scan'])
def scan(message):
	tfm_tb.reply_to(message, "Scanning...")
	log.info("Scanning...")
	scan_dir = "/srv/share/hfs/misc/tfm_temp/scan"
	files = []
	for file in os.listdir(scan_dir):
		new_file = {"name": file}
		file = os.path.join(scan_dir, file)
		if not os.path.isfile(file):
			continue
		new_file["path"] = file
		try:
			exif = ljson(os.popen(f"exiftool -json \"{file}\"").read())[0]
		except Exception as e:
			log.error("Error while parsing EXIF for file '%s': %s" % (file, e))
			continue
		dt = exif["FileModifyDate"]
		if "SubSecDateTimeOriginal" in exif.keys():
			dt = exif["SubSecDateTimeOriginal"]
		elif "DateTimeOriginal" in exif.keys():
			dt = exif["DateTimeOriginal"]
		if '.' in dt:
			try:
				dt = datetime.strptime(dt, "%Y:%m:%d %H:%M:%S.%f%z")
			except:
				dt = TZ.localize(datetime.strptime(dt, "%Y:%m:%d %H:%M:%S.%f"))
		else:
			try:
				try:
					dt = datetime.strptime(dt, "%Y:%m:%d %H:%M:%S%z")
				except:
					dt = TZ.localize(datetime.strptime(dt[:19], "%Y:%m:%d %H:%M:%S"))
			except:
				log.error("Broken date: %s\t%s" % (new_file, dt))
				continue
		new_file["datetime"] = dt
		files.append(new_file)
	tfm_tb.reply_to(message, f"{len(files)} files found.")
	log.info(f"{len(files)} files found.")
	files.sort(key=lambda f: f["datetime"])
	for file in files:
		try:
			file_id, ext = tfm.add_file(file["path"], file["datetime"])
			tfm.add_file_to_tag(file_id, "d6d8129a-984d-4451-8c83-d04523ced8a8")
		except Exception as e:
			tfm_tb.reply_to(message, f"Error adding file '{file['name']}': {str(e)}")
			log.error(f"Error adding file '{file['name']}': {str(e)}")
	tfm_tb.reply_to(message, "Files added.")
	log.info("Files added.")


if __name__ == '__main__':
	log.info("tfm-tb started")
	defer("log.info(\"tfm-tb stopped\")")
	while True:
		try:
			tfm_tb.polling(none_stop=True)
		except Exception as e:
			log.exception("exception on polling")
			sleep(1)