Image info, endpoint encryption

bot/db/__init__.py

@@ -1,2 +1,3 @@
 from .db import db
 from .db_model import DBTables
+from .encryption import encrypt, decrypt

bot/db/db.py

@@ -11,5 +11,6 @@ db = {
     'cooldown': DBDict(DB, autocommit=True, tablename='cooldown'),
     'exceptions': DBDict(DB, autocommit=True, tablename='exceptions'),
     'queue': DBDict(DB, autocommit=True, tablename='queue'),
-    'generated': DBDict(DB, autocommit=True, tablename='generated')
+    'generated': DBDict(DB, autocommit=True, tablename='generated'),
+    'prompts': DBDict(DB, autocommit=True, tablename='prompts')
 }

bot/db/db_model.py

@@ -8,12 +8,13 @@ from .meta import DBMeta
 
 
 class DBTables:
-    tables = ['config', 'cooldown', 'exceptions', 'queue', 'generated']
+    tables = ['config', 'cooldown', 'exceptions', 'queue', 'generated', 'prompts']
     config = "config"
     cooldown = "cooldown"
     exceptions = "exceptions"
     queue = "queue"
     generated = "generated"
+    prompts = "prompts"
 
 
 class DBDict(SqliteDict):

bot/db/encryption.py

@@ -0,0 +1,25 @@
+from cryptography.fernet import Fernet
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+import base64
+from bot.config import ENCRYPTION_KEY, BARS_APP_ID
+
+
+fernet = Fernet(
+    base64.urlsafe_b64encode(
+        PBKDF2HMAC(
+            algorithm=hashes.SHA256(),
+            length=32,
+            iterations=390000,
+            salt=BARS_APP_ID.encode()
+        ).derive(ENCRYPTION_KEY)
+    )
+)
+
+
+def encrypt(s: str) -> bytes:
+    return fernet.encrypt(s.encode())
+
+
+def decrypt(s: bytes) -> str:
+    return fernet.decrypt(s).decode()