80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
"""
|
|
models/db.py — DBfrontend: encrypted message store.
|
|
"""
|
|
|
|
import time
|
|
from typing import Optional
|
|
from cryptography.fernet import Fernet
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
|
import base64
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from .base import make_engine, Messages
|
|
|
|
|
|
class DBfrontend:
|
|
def __init__(self, db_cfg, secret: bytes, salt: bytes):
|
|
self.engine = make_engine(db_cfg)
|
|
self.Session = sessionmaker(bind=self.engine)
|
|
self.cipher = self._get_cipher(secret, salt)
|
|
|
|
@staticmethod
|
|
def _get_cipher(secret: bytes, salt: bytes) -> Fernet:
|
|
kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100_000)
|
|
key = base64.urlsafe_b64encode(kdf.derive(secret))
|
|
return Fernet(key)
|
|
|
|
def _encrypt(self, message: str) -> str:
|
|
return self.cipher.encrypt(message.encode()).decode()
|
|
|
|
def _decrypt(self, encrypted: str) -> str:
|
|
try:
|
|
return self.cipher.decrypt(encrypted.encode()).decode()
|
|
except Exception:
|
|
return "[Decryption failed]"
|
|
|
|
def add_message(self, channel: str, ip: str, content: str, gif_url: Optional[str] = None) -> int:
|
|
with self.Session() as session:
|
|
msg = Messages(
|
|
channel=channel,
|
|
ip=ip,
|
|
content=self._encrypt(content),
|
|
gif_url=gif_url,
|
|
timestamp=int(time.time() * 1000),
|
|
)
|
|
session.add(msg)
|
|
session.commit()
|
|
return msg.id
|
|
|
|
def get_messages(self, channel: str = "general", limit: int = 50) -> list:
|
|
with self.Session() as session:
|
|
q = (
|
|
session.query(Messages)
|
|
.filter(Messages.channel == channel)
|
|
.order_by(Messages.id.desc())
|
|
.limit(limit)
|
|
)
|
|
result = []
|
|
for msg in reversed(q.all()):
|
|
result.append(
|
|
{
|
|
"id": msg.id,
|
|
"channel": msg.channel,
|
|
"ip": msg.ip,
|
|
"content": self._decrypt(msg.content),
|
|
"gif_url": msg.gif_url,
|
|
"timestamp": msg.timestamp,
|
|
}
|
|
)
|
|
return result
|
|
|
|
def delete_message(self, msg_id: int) -> bool:
|
|
with self.Session() as session:
|
|
msg = session.get(Messages, msg_id)
|
|
if msg:
|
|
session.delete(msg)
|
|
session.commit()
|
|
return True
|
|
return False
|