libmfs.py

import os
import io
import gzip
import json
import uuid
import math
import random
import datetime
import requests
import humanize
import mimetypes
from dataclasses import dataclass
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Hash import SHA256

DEFAULT_MAX_FILE_SIZE = 10 * 1024 * 1024
DEFAULT_FETCH_LIMIT = 1000

class FileNotFoundInRoom(Exception):
    pass

class AESCipher:
    def __init__(self, key: str):
        h = SHA256.new()
        h.update(key.encode())
        self.key = h.digest()

    def encrypt(self, raw: bytes) -> bytes:
        cipher = AES.new(self.key, AES.MODE_CBC)
        ct_bytes = cipher.encrypt(pad(raw, AES.block_size))
        return cipher.iv + ct_bytes

    def decrypt(self, enc: bytes) -> bytes:
        iv = enc[:16]
        ct = enc[16:]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        return unpad(cipher.decrypt(ct), AES.block_size)

    def encrypt_text(self, text: str) -> str:
        return self.encrypt(text.encode("utf-8")).hex()

    def decrypt_text(self, text: str) -> str:
        raw = bytes.fromhex(text)
        return self.decrypt(raw).decode("utf-8")

@dataclass
class MatrixMetadata:
    size: int
    part: int
    total_part: int
    date_added: float
    filename: str

    def to_json(self) -> str:
        return json.dumps(self.__dict__)

    @staticmethod
    def from_json(text: str):
        return MatrixMetadata(**json.loads(text))

class MatrixFs:
    def __init__(self, server_url, username, password, room_name, encryption_key,
                 max_file_size=DEFAULT_MAX_FILE_SIZE, fetch_limit=DEFAULT_FETCH_LIMIT, variance=0.10):
        self.server_url = server_url.rstrip("/")
        self.username = username
        self.password = password
        self.room_name = room_name
        self.max_file_size = max_file_size
        self.fetch_limit = fetch_limit
        self.verify_ssl = False
        self.token_file = f".token_{username}"
        self.token = None
        self.user_id = None
        self.room_id = None
        self.cipher = AESCipher(encryption_key)
        self.variance = variance

        self.token = self._load_token()
        self.user_id = self._check_token() if self.token else None
        if not self.user_id:
            self.user_id = self._login()
        self.room_id = self._find_or_create_room()

    def generate_filename(self):
        subjects = [
            "photo", "image", "selfie", "dessin", "liste", "truc",
            "bidule", "recette", "courses", "vidéo", "chanson",
            "copie", "vieux_fichier", "idée", "notes"
        ]
        
        topics = [
            "chlamydia","syphilis","gonorrhée","herpès","papillomavirus","urétrite","condylome","chancroïde","vaginose","candidose","infection_urinaire","prurit","douleurs_pelviennes","pertes_bizarres","boutons","odeur_forte","test_positif","écoulement",
            "murielle", "chat", "chien", "noel", "vacances", "voiture", "canapé",
            "fromage", "tartiflette", "lasagne", "pâtes", "fondue", "gâteau",
            "pizza", "croquettes", "quiche", "gratin", "bouffe", "placard",
            "jardin", "verrue", "mycose", "hemoroides", "impuissance", "calvitie",
            "tripes", "herpes", "chaude_pisse","furoncle", "orteils", "chaussettes", "kebab", "rot", "flatulences",
            "Sophie","Julie","Chloé","Sarah","Marine","Lucie","Laura","Manon",
            "Clara","Océane","Élodie","Léa","Mélanie","Amandine","Justine","Pauline","Caroline",
            "Céline","Anaïs","Audrey","Nathalie","Isabelle","Vanessa","Jessica","Marion","Amélie",
            "Sandra","Valérie", "Thomas","Lucas","Julien","Maxime","Alexandre","Nicolas","Antoine","Hugo","Kevin","Matthieu",
            "Clément","Romain","Pierre","Vincent","Florian","Guillaume","Benoît","Quentin","Théo",
            "Benjamin","Sébastien","Jérôme","Aurélien","Loïc","Cédric","Paul","Samuel","Mickaël",
            "Jean"
        ]

        suffixes = ["new","new2","nouveau", "v2", "final", "copie", "def", "old", "old2", "2", "4","dernier","test"]
        extensions = ["pdf", "docx", "txt", "jpg", "png", "zip", "mp4", "avi", "xls","gif","avi","7z","bmp"]

        subject = random.choice(subjects)
        topic = random.choice(topics)
        ext = random.choice(extensions)
        seps = [" ", ".","-","_"]
        suffix = random.choice(seps) + random.choice(suffixes) if random.random() < 0.4 else ""
        
        return f"{subject}_{topic}{suffix}.{ext}"

    def _save_token(self, token):
        with open(self.token_file, "w") as f:
            f.write(token)

    def _load_token(self):
        if os.path.exists(self.token_file):
            with open(self.token_file) as f:
                return f.read().strip()
        return None

    def _check_token(self):
        try:
            r = self._do_request("GET", "/_matrix/client/v3/account/whoami")
            return r.json()["user_id"]
        except:
            return None

    def _login(self):
        payload = {
            "type": "m.login.password",
            "user": self.username,
            "password": self.password
        }
        r = self._do_request("POST", "/_matrix/client/v3/login", json=payload)
        data = r.json()
        self.token = data["access_token"]
        self._save_token(self.token)
        return data["user_id"]

    def _build_url(self, endpoint: str) -> str:
        return f"{self.server_url}{endpoint}"

    def _do_request(self, method, endpoint, **kwargs):
        headers = kwargs.get("headers", {})
        if self.token:
            headers["Authorization"] = f"Bearer {self.token}"
        kwargs["headers"] = headers
        url = self._build_url(endpoint)
        response = requests.request(method, url, verify=self.verify_ssl, **kwargs)
        response.raise_for_status()
        return response

    def _get_joined_rooms(self):
        return self._do_request("GET", "/_matrix/client/v3/joined_rooms").json()["joined_rooms"]

    def _find_or_create_room(self):
        for room_id in self._get_joined_rooms():
            try:
                r = self._do_request("GET", f"/_matrix/client/v3/rooms/{room_id}/state/m.room.name")
                if r.json().get("name") == self.room_name:
                    return room_id
            except:
                continue

        payload = {
            "name": self.room_name,
            "preset": "private_chat",
            "visibility": "private",
            "initial_state": [{
                "type": "m.room.encryption",
                "state_key": "",
                "content": {"algorithm": "m.megolm.v1.aes-sha2"}
            }]
        }
        r = self._do_request("POST", "/_matrix/client/v3/createRoom", json=payload)
        return r.json()["room_id"]

    def _upload_bytes(self, data: bytes, filename: str, mime: str) -> str:
        headers = {"Content-Type": mime}
        params = {"filename": filename}
        r = self._do_request("POST", "/_matrix/media/v3/upload", headers=headers, params=params, data=data)
        return r.json()["content_uri"]

    def _post_message(self, mxc_url: str, fake_name: str, encrypted_chunk: bytes, original_filename: str, part_index: int, total_parts: int, total_size: int):
        metadata = MatrixMetadata(
            size=total_size,
            part=part_index,
            total_part=total_parts,
            date_added=datetime.datetime.now().timestamp(),
            filename=original_filename
        )
        payload = {
            "msgtype": "m.file",
            "body": "",
            "filename": fake_name,
            "url": mxc_url,
            "info": {
                "mimetype": mimetypes.types_map.get(os.path.splitext(fake_name)[1].lower()) or "application/octet-stream",
                "size": len(encrypted_chunk)
            },
            "text": self.cipher.encrypt_text(metadata.to_json()),
        }
        txn = str(uuid.uuid4())
        self._do_request(
            "PUT",
            f"/_matrix/client/v3/rooms/{self.room_id}/send/m.room.message/{txn}",
            json=payload
        )

    def send_file(self, filepath: str):
        filename = os.path.basename(filepath)

        with open(filepath, "rb") as f:
            raw_data = f.read()

        gzip_buffer = io.BytesIO()
        with gzip.GzipFile(fileobj=gzip_buffer, mode="wb") as gz:
            gz.write(raw_data)
        data = gzip_buffer.getvalue()

        total_size = len(data)
        print(f"Compressed {humanize.naturalsize(len(raw_data), binary=False)} -> {humanize.naturalsize(len(data), binary=False)}")
        mime = "application/octet-stream"

        base_chunk_size = self.max_file_size
        max_parts = math.ceil(total_size / base_chunk_size)

        chunk_sizes = []
        remaining = total_size
        for i in range(max_parts - 1):
            variance = random.uniform(-self.variance, self.variance)
            chunk_size = int(base_chunk_size * (1 + variance))
            chunk_size = min(chunk_size, remaining)
            chunk_sizes.append(chunk_size)
            remaining -= chunk_size

        chunk_sizes.append(remaining)
        total_parts = len(chunk_sizes)

        offset = 0
        for i, chunk_size in enumerate(chunk_sizes):
            chunk = data[offset:offset + chunk_size]
            offset += chunk_size

            encrypted_chunk = self.cipher.encrypt(chunk)
            mxc_url = self._upload_bytes(encrypted_chunk, filename, mime)
            fake_name = self.generate_filename()
            self._post_message(mxc_url, fake_name, encrypted_chunk, filename, i + 1, total_parts, total_size)

        print(f"File '{filename}' uploaded in {total_parts} part(s).")

    def _download_file(self, mxc_url):
        if not mxc_url.startswith("mxc://"):
            raise ValueError("Invalid MXC URL")
        server, media_id = mxc_url[6:].split("/", 1)
        r = self._do_request("GET", f"/_matrix/client/v1/media/download/{server}/{media_id}")
        return r.content

    def find_file(self, filename: str) -> bytes:
        messages = self._fetch_room_messages()
        parts = {}

        for event in messages:
            content = event.get("content", {})
            if content.get("msgtype") != "m.file":
                continue

            try:
                metadata = MatrixMetadata.from_json(self.cipher.decrypt_text(content["text"]))
                if metadata.filename != filename:
                    continue
                mxc_url = content["url"]
                encrypted_data = self._download_file(mxc_url)
                parts[metadata.part] = self.cipher.decrypt(encrypted_data)
            except Exception as e:
                print(f"Error processing file part: {e}")

        if not parts:
            raise FileNotFoundInRoom(f"File '{filename}' not found in room '{self.room_name}'.")

        expected_parts = list(range(1, max(parts.keys()) + 1))
        if sorted(parts.keys()) != expected_parts:
            raise FileNotFoundInRoom(f"Incomplete file: missing parts for '{filename}'.")

        compressed_data = b"".join(parts[i] for i in expected_parts)

        with gzip.GzipFile(fileobj=io.BytesIO(compressed_data), mode="rb") as gz:
            decompressed_data = gz.read()

        return decompressed_data

    def delete(self, filename: str):
        messages = self._fetch_room_messages()
        found = False

        for event in messages:
            content = event.get("content", {})
            if content.get("msgtype") != "m.file":
                continue

            try:
                meta = MatrixMetadata.from_json(self.cipher.decrypt_text(content["text"]))
                if meta.filename != filename:
                    continue

                event_id = event.get("event_id")
                if event_id:
                    txn = str(uuid.uuid4())
                    self._do_request("PUT", f"/_matrix/client/v3/rooms/{self.room_id}/redact/{event_id}/{txn}", json={})
                    found = True
            except Exception as e:
                print(f"Error while deleting message: {e}")

        if not found:
            raise FileNotFoundInRoom(f"File '{filename}' not found for deletion.")

    def list_files(self):
        messages = self._fetch_room_messages()
        files = {}
        for event in messages:
            content = event.get("content", {})
            if content.get("msgtype") != "m.file":
                continue
            try:
                meta = MatrixMetadata.from_json(self.cipher.decrypt_text(content["text"]))
                name = meta.filename
                _, media_id = content["url"][6:].split("/", 1)
                if name not in files:
                    files[name] = {
                        "size": meta.size,
                        "chunks": meta.total_part,
                        "date_added": meta.date_added,
                        "media_id": media_id,
                   }
                else:
                    files[name]["size"] = max(files[name]["size"], meta.size)
                    files[name]["chunks"] = max(files[name]["chunks"], meta.total_part)
            except Exception as e:
                print(f"Error in list_files: {e}")
                continue

        return files

    def _fetch_room_messages(self):
        try:
            r = self._do_request(
                "GET",
                f"/_matrix/client/v3/rooms/{self.room_id}/messages",
                params={"limit": self.fetch_limit, "dir": "b"},
            )
            return r.json().get("chunk", [])
        except Exception as e:
            print(f"Error fetching messages: {e}")
            return []