tracker.py

import socket
import threading
import json
import os
from datetime import datetime, timezone
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import serialization, hashes

HOST = 'localhost'
PORT = 8000

USERS_FILE = "users.json"
ONLINE_PEERS_FILE = "online_peers.json"
CHAT_ROOMS_FILE = "chat_rooms.json"
CHAT_MESSAGES_FILE = "chat_messages.json"
ROOM_ADMINS_FILE = "room_admins.json"
PRIVATE_KEY_FILE = "private_key.pem"

IP_BASE = "127.0.0.1"  # IP fixo localhost para testes locais
PORT_POOL_START = 50000
PORT_POOL_END = 60000

users_db = {}
online_peers = {}
chat_rooms = {}
chat_messages = {}
room_admins = {}

allocated_ports = set()

lock = threading.Lock()

def load_or_create_private_key():
    if os.path.exists(PRIVATE_KEY_FILE):
        with open(PRIVATE_KEY_FILE, "rb") as f:
            private_key = serialization.load_pem_private_key(f.read(), password=None)
        print("[i] Chave privada carregada do arquivo.")
    else:
        private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        with open(PRIVATE_KEY_FILE, "wb") as f:
            pem = private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption())
            f.write(pem)
        print("[i] Nova chave privada gerada e salva.")
    return private_key

private_key = load_or_create_private_key()
public_key = private_key.public_key()

public_key_pem = public_key.public_bytes(
    encoding=serialization.Encoding.PEM,
    format=serialization.PublicFormat.SubjectPublicKeyInfo
).decode()


def save_json(filename, data):
    plaintext = json.dumps(data, indent=2).encode("utf-8")
    chunks = [plaintext[i:i + 190] for i in range(0, len(plaintext), 190)]
    encrypted_chunks = [public_key.encrypt(chunk,
        padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)) for chunk in chunks]
    with open(filename, "wb") as f:
        for chunk in encrypted_chunks:
            f.write(len(chunk).to_bytes(4, byteorder="big") + chunk)


def load_json(filename):
    if not os.path.exists(filename):
        return {}
    with open(filename, "rb") as f:
        encrypted_data = f.read()
    decrypted_chunks = []
    i = 0
    try:
        while i < len(encrypted_data):
            chunk_size = int.from_bytes(encrypted_data[i:i + 4], byteorder="big")
            i += 4
            encrypted_chunk = encrypted_data[i:i + chunk_size]
            i += chunk_size
            decrypted_chunk = private_key.decrypt(encrypted_chunk,
                padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None))
            decrypted_chunks.append(decrypted_chunk)
    except Exception as e:
        print(f"[!] Erro ao descriptografar arquivo {filename}: {e}")
        print("[!] O arquivo pode estar corrompido ou não criptografado corretamente.")
        return {}
    plaintext = b''.join(decrypted_chunks).decode("utf-8")
    return json.loads(plaintext)


def load_all():
    global users_db, online_peers, chat_rooms, chat_messages, allocated_ports, room_admins
    users_db = load_json(USERS_FILE)
    online_peers = load_json(ONLINE_PEERS_FILE)
    chat_rooms = load_json(CHAT_ROOMS_FILE)
    chat_messages = load_json(CHAT_MESSAGES_FILE)
    room_admins = load_json(ROOM_ADMINS_FILE)
    allocated_ports = set(peer["port"] for peer in online_peers.values())


def save_all():
    save_json(USERS_FILE, users_db)
    save_json(ONLINE_PEERS_FILE, online_peers)
    save_json(CHAT_ROOMS_FILE, chat_rooms)
    save_json(CHAT_MESSAGES_FILE, chat_messages)
    save_json(ROOM_ADMINS_FILE, room_admins)


def allocate_ip_port():
    for port_candidate in range(PORT_POOL_START, PORT_POOL_END + 1):
        if port_candidate not in allocated_ports:
            allocated_ports.add(port_candidate)
            return IP_BASE, port_candidate
    return None, None


def release_ip_port(ip, port):
    allocated_ports.discard(port)


def decrypt_message(encrypted_data):
    return private_key.decrypt(
        encrypted_data,
        padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
    ).decode()


def handle_client(conn, addr):
    print(f"[+] Nova conexão de {addr}")
    username = None
    try:
        conn.sendall(public_key_pem.encode())
        data = conn.recv(4096)
        if not data:
            return
        msg_json = decrypt_message(data)
        message = json.loads(msg_json)
        response = handle_message(message, addr)
        if response.get("status") == "ok" and message.get("type") == "login":
            username = message.get("username")
            allocated_ip = online_peers[username]["ip"]
            allocated_port = online_peers[username]["port"]
        conn.sendall(json.dumps(response).encode())
    except Exception as e:
        print(f"[!] Erro ao lidar com {addr}: {e}")
    finally:
        conn.close()
        print(f"[-] Conexão encerrada de {addr}")


def handle_message(msg, addr):
    msg_type = msg.get("type")
    username = msg.get("username")
    password = msg.get("password")

    with lock:
        if msg_type == "register":
            if username in users_db:
                return {"status": "error", "message": "Usuário já existe"}
            users_db[username] = password
            save_json(USERS_FILE, users_db)
            return {"status": "ok", "message": "Registrado com sucesso"}

        elif msg_type == "login":
            if users_db.get(username) != password:
                return {"status": "error", "message": "Credenciais inválidas"}
            if username in online_peers:
                peer_info = online_peers[username]
            else:
                ip, port = allocate_ip_port()
                if ip is None or port is None:
                    return {"status": "error", "message": "Não há porta disponível"}
                peer_info = {"ip": ip, "port": port}
                online_peers[username] = peer_info
                save_json(ONLINE_PEERS_FILE, online_peers)
            return {"status": "ok", "message": "Login bem-sucedido", "ip": peer_info["ip"], "port": peer_info["port"]}

        elif msg_type == "logout":
            if username in online_peers:
                ip = online_peers[username]["ip"]
                port = online_peers[username]["port"]
                release_ip_port(ip, port)
                del online_peers[username]
                save_json(ONLINE_PEERS_FILE, online_peers)
                return {"status": "ok", "message": "Logout efetuado"}
            else:
                return {"status": "error", "message": "Usuário não está online"}

        elif msg_type == "list_peers":
            return {"status": "ok", "peers": online_peers}

        elif msg_type == "create_room":
            room = msg.get("room")
            if room in chat_rooms:
                return {"status": "error", "message": "Sala já existe"}
            chat_rooms[room] = [username]
            chat_messages[room] = []
            room_admins[room] = username
            save_all()
            return {"status": "ok", "message": f"Sala '{room}' criada"}

        elif msg_type == "join_room":
            room = msg.get("room")
            if room not in chat_rooms:
                return {"status": "error", "message": "Sala não existe"}

            admin = room_admins.get(room)
            membros = chat_rooms[room]

            if username == admin or username in membros:
                return {"status": "ok", "message": f"Entrou na sala '{room}'"}
            else:
                return {"status": "error", "message": "Acesso negado. Você não foi convidado para esta sala."}


        elif msg_type == "invite_to_room":
            room = msg.get("room")
            target_user = msg.get("target")
            if room not in chat_rooms:
                return {"status": "error", "message": "Sala não existe"}
            if room_admins.get(room) != username:
                return {"status": "error", "message": "Apenas o administrador pode convidar"}
            if target_user not in online_peers:
                return {"status": "error", "message": "Peer alvo offline"}
            if target_user not in chat_rooms[room]:
                chat_rooms[room].append(target_user)
                save_json(CHAT_ROOMS_FILE, chat_rooms)
            return {"status": "ok", "message": f"{target_user} convidado para sala '{room}'"}

        elif msg_type == "remove_from_room":
            room = msg.get("room")
            target_user = msg.get("target")
            if room not in chat_rooms:
                return {"status": "error", "message": "Sala não existe"}
            if room_admins.get(room) != username:
                return {"status": "error", "message": "Apenas o administrador pode remover"}
            if target_user in chat_rooms[room]:
                chat_rooms[room].remove(target_user)
                save_json(CHAT_ROOMS_FILE, chat_rooms)
                return {"status": "ok", "message": f"{target_user} removido da sala '{room}'"}
            else:
                return {"status": "error", "message": "Usuário não está na sala"}

        elif msg_type == "list_rooms":
            return {"status": "ok", "rooms": list(chat_rooms.keys())}

        elif msg_type == "get_room_members":
            room = msg.get("room")
            if room not in chat_rooms:
                return {"status": "error", "message": "Sala não existe"}
            return {"status": "ok", "members": chat_rooms[room]}

        elif msg_type == "get_room_admin":
            room = msg.get("room")
            if room not in room_admins:
                return {"status": "error", "message": "Sala sem administrador definido"}
            return {"status": "ok", "admin": room_admins[room]}

        elif msg_type == "send_message_room":
            room = msg.get("room")
            message_text = msg.get("message")
            if room not in chat_rooms:
                return {"status": "error", "message": "Sala não existe"}
            if username not in chat_rooms[room]:
                return {"status": "error", "message": "Você não está na sala"}
            msg_obj = {
                "username": username,
                "message": message_text,
                "timestamp": datetime.now(timezone.utc).isoformat()
            }
            if room not in chat_messages:
                chat_messages[room] = []
            chat_messages[room].append(msg_obj)
            save_json(CHAT_MESSAGES_FILE, chat_messages)
            return {"status": "ok", "message": f"{username}: \"{message_text}\""}

        elif msg_type == "get_room_messages":
            room = msg.get("room")
            if room not in chat_messages:
                return {"status": "error", "message": "Sala não existe ou sem mensagens"}
            return {"status": "ok", "messages": chat_messages[room]}

        else:
            return {"status": "error", "message": "Tipo de mensagem inválido"}


def start_tracker():
    load_all()
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind((HOST, PORT))
    server.listen()
    print(f"[i] Tracker iniciado em {HOST}:{PORT}")
    while True:
        conn, addr = server.accept()
        thread = threading.Thread(target=handle_client, args=(conn, addr))
        thread.daemon = True
        thread.start()


if __name__ == "__main__":
    start_tracker()