diff --git a/assistant/config.py b/assistant/config.py index fa79d0e..651b298 100644 --- a/assistant/config.py +++ b/assistant/config.py @@ -22,6 +22,9 @@ class Config: API_HASH = os.environ.get("API_HASH") BOT_TOKEN = os.environ.get("BOT_TOKEN") AUTH_CHATS = set([-1001481357570]) # @UserGeOt + AES_DB_PASS = str(os.environ.get("AES_DB_PASS")) + DB_VI_KEY = str(os.environ.get("DB_VI_KEY")) + DB_CHANNEL = int(os.environ.get("DB_CHANNEL")) if os.environ.get("AUTH_CHATS"): AUTH_CHATS.update(map(int, os.environ.get("AUTH_CHATS").split())) WHITELIST_CHATS = set([-1001465749479]) # @UserGeSpam diff --git a/assistant/j_ses.py b/assistant/j_ses.py new file mode 100644 index 0000000..afe38de --- /dev/null +++ b/assistant/j_ses.py @@ -0,0 +1,124 @@ +''' JSON session manager ''' +import os +import json +from hashlib import blake2s + +import aiofiles +from PIL import Image +from pySmartDL import SmartDL + +from pyrogram import Client, types, crypto + +from assistant import Config + +AES_PASS = blake2s(Config.AES_DB_PASS.encode()).digest() +VI_KEY = blake2s(Config.DB_VI_KEY.encode()).digest() +LOG_ID = Config.DB_CHANNEL + +SES_DIR = './.json-session' +SES_FILZ = { + 'test': 'test-session.json' +} + + +class JsonSeVe: + ''' ╮(╯▽╰)╭ ''' + + def __init__(self, session: str, save_id: int = 0): + ''' Initialization of Json Save + param: session :-> name of session like warn, etc that exists in SES_FILZ + param: save_id :-> message id of a media that is to be edited. + ''' + self.jthumb = os.path.abspath('./.thumbs/jsumb.jpeg') + self.thumb_uri = 'https://telegra.ph/file/da08e43d19d5023a248be.jpg' + self.tg_save = save_id + self.dir = SES_DIR + self.session = session + self.file = SES_FILZ.get(self.session) + self.ses_path = os.path.abspath(str(self.dir) + '/' + (self.file)) + self.ses_path_enc = self.ses_path + '.enc' + if not os.path.lexists(self.dir): + os.makedirs(self.dir) + if not os.path.lexists('./.thumbs'): + os.makedirs('./.thumbs') + if not os.path.isfile(self.jthumb): + dl_thumb = SmartDL(self.thumb_uri, self.jthumb, progress_bar=True) + dl_thumb.start(blocking=True) + img = Image.open(self.jthumb) + img.resize((320, 320)) + img.save(self.jthumb) + + @staticmethod + def _add_buffer(data: bytes): + ''' add white spaces as multiples of 16 is required for encryption ''' + lez = len(data) + if lez % 16 != 0: + data += b' ' * (16 - lez % 16) + return data + + async def encrypt_data(self): + ''' makes a copy of encrypted data ''' + async with aiofiles.open(self.ses_path, 'rb') as enc: + data = await enc.readlines() + enc_data = [] + for line in data: + line = self._add_buffer(line) + enc_data.append(crypto.aes.ige256_encrypt(line, AES_PASS, VI_KEY)) + async with aiofiles.open(self.ses_path_enc, 'wb') as wer: + await wer.writelines(enc_data) + + async def decrypt_data(self): + ''' decrpts encrypted file ''' + async with aiofiles.open(self.ses_path_enc, 'rb') as dec: + data = await dec.readlines() + dec_data = [] + for line in data: + dec_data.append(crypto.aes.ige256_decrypt(line, AES_PASS, VI_KEY).strip()) + async with aiofiles.open(self.ses_path, 'wb') as wer: + await wer.writelines(dec_data) + + async def load_session(self, client: Client): + ''' loadz session from telegram ''' + save: types.Message = await client.get_messages(LOG_ID, self.tg_save) + doc = save.document + if doc and doc.file_name == self.file + '.enc': + await client.download_media( + message=save, + file_name=self.ses_path_enc, + block=True, + ) + await self.decrypt_data() + return True + return False + + async def save_session(self, client: Client): + ''' exports file to telegram ''' + if not os.path.isfile(self.ses_path): + # Try a Logging + return False + await self.encrypt_data() + await client.edit_message_media( + chat_id=LOG_ID, + message_id=self.tg_save, + media=types.InputMediaDocument( + media=self.ses_path_enc, + thumb=self.jthumb + ) + ) + return True + + async def get_json(self, client: Client): + ''' returns json data in a file ''' + if not os.path.isfile(self.ses_path): + is_save = await self.load_session(client) + if not is_save: + return {} + with open(self.ses_path, 'rb') as j_s: + data = json.load(j_s) + return data + + async def write_json(self, client: Client, data: dict): + ''' writes a json data to file [this may flush entire file]''' + with open(self.ses_path, 'w') as j_s: + json.dump(data, j_s) + return await self.save_session(client) diff --git a/requirements.txt b/requirements.txt index 5d109d7..1c31fb4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,7 @@ aiofiles aiohttp +Pillow pyrogram +pySmartDL python-dotenv +tgcrypto