chatserver/ChatServer/CS3.1.py

import threading
import json
from flask import Flask, jsonify, request
import sqlite3
import socket
import base64
import secrets
import time
app = Flask(__name__)
socket_server = socket.socket()
socket_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

active_users = {}
chat_connections = []
tokens = {}

def get_db_connection():
    conn = sqlite3.connect("usr.db")
    conn.row_factory = sqlite3.Row
    return conn

def isuserxist(name):
    cn = get_db_connection()
    csr = cn.cursor()
    csr.execute('SELECT * FROM users WHERE name = ?', (name,))
    rst = csr.fetchone()
    cn.close()
    if rst is not None:
        return True
    else:
        return False

def ispsswdright(name,passwd):
    cn = get_db_connection()
    csr = cn.cursor()
    csr.execute("SELECT COUNT(*) FROM users WHERE name=?", (name,))
    row_count = csr.fetchone()[0]
    password = None
    if row_count > 0:
        csr.execute("SELECT passwd FROM users WHERE name=?", (name,))
        password = csr.fetchone()[0]
    if password == passwd:
        return True
    else:
        return False

def register_user(usr, pwd):
    conn = get_db_connection()
    csr2 = conn.cursor()
    csr2.execute('SELECT * FROM users WHERE name = ?', (usr,))
    result = csr2.fetchone()
    if result is not None:
        return {"type": "register_0", "success": False, "message": "Username already exists"}
    else:
        try:
            cursor = conn.cursor()
            cursor.execute("INSERT INTO users (name, passwd) VALUES (?, ?)", (usr, pwd))
            conn.commit()
            return {"type": "register_1", "success": True, "message": "User registered successfully"}
        except sqlite3.Error as e:
            return {"type": "register_0", "success": False, "message": str(e)}
        finally:
            conn.close()

def generate_token(username):
    token = secrets.token_hex(16)
    tokens[token] = {'username': username, 'timestamp': time.time()}
    return token

def validate_token(token):
    if token in tokens:
        if time.time() - tokens[token]['timestamp'] < 3600:
            tokens[token]['timestamp'] = time.time()
            return tokens[token]['username']
    return None

@app.route("/api/register", methods=['POST'])
def register1():
    vl = request.get_json()
    usr = vl.get('username')
    pwd = vl.get('password')
    result = register_user(usr, pwd)
    if result['success']:
        return jsonify(result)
    else:
        return jsonify(result), 403 if result['message'] == "Username already exists" else 500

@app.route("/api/login", methods=['POST'])
def login():
    data = request.get_json()
    if isuserxist(data['username']):
        if ispsswdright(data['username'], data['password']):
            token = generate_token(data['username'])
            return jsonify({"type": "login_1", "status": "success", "token": token})
    return jsonify({"type": "login_0", "status": "error"}), 401

@app.route("/api/chat", methods=['POST'])
def chat():
    token = request.headers.get('Authorization')
    username = validate_token(token)
    if not username:
        return jsonify({"type": "chat", "status": "error"}), 401
    data = request.get_json()
    message = {
        "type": "chat",
        "user": username,
        "message": data['message']
    }
    broadcast_message(message)
    return jsonify({"type": "chat", "status": "success"})

def broadcast_message(message, sender=None):
    for conn in chat_connections:
        try:
            conn.sendall(json.dumps(message).encode())
        except:
            chat_connections.remove(conn)

def handle_socket_message(data, addr, conn):
    try:
        action = data.get('type')
        if action == 'register':
            result = register_user(data.get('username'), data.get('password'))
            if result['success']:
                return {"type": "register_1","status": "success", "message": result['message']}
            else:
                return {"type": "register_0","status": "error", "message": result['message']}

        elif action == 'login':
            if isuserxist(data['username']):
                if ispsswdright(data['username'], data['password']):
                    active_users[addr[0]] = data['username']
                    tk = base64.b64encode(data['username'].encode('utf-8'))
                    chat_connections.append(conn)
                    return {"type": "login_1", "status": "success", "message": "Login successful", "token": generate_token(data['username']), "username": data['username']}
            return {"type": "login_0", "status": "error", "message": "Invalid credentials", "username": data['username']}
        elif action == 'chat':
            if addr[0] in active_users:
                message = {
                    "type": "chat",
                    "user": active_users[addr[0]],
                    "message": data['message']
                }
                broadcast_message(message)
                return {"type": "chat", "status": "success"}
            return {"type": "chat", "status": "error", "message": "Not logged in"}
    except Exception as e:
        return {"status": "error", "message": str(e)}

def run_socket_server():
    socket_server.bind(("localhost", 8889))
    socket_server.listen()
    print("Socket server running on port 8889")
    while True:
        conn, addr = socket_server.accept()
        print(f"Socket client connected: {addr}")
        try:
            while True:
                data = conn.recv(1024)
                if not data:
                    if addr[0] in active_users:
                        del active_users[addr[0]]
                    break

                try:
                    json_data = json.loads(data.decode())
                    response = handle_socket_message(json_data, addr, conn)
                    conn.sendall(json.dumps(response).encode())
                except json.JSONDecodeError:
                    conn.sendall(json.dumps(
                        {"type": "register_0", "status": "error", "message": "Invalid JSON"}
                    ).encode())
        except ConnectionResetError:
            if addr[0] in active_users:
                del active_users[addr[0]]
            print(f"Client {addr} disconnected")
        finally:
            if conn in chat_connections:
                chat_connections.remove(conn)
            conn.close()

if __name__ == '__main__':
    with get_db_connection() as conn:
        conn.execute('''CREATE TABLE IF NOT EXISTS users
                      (name TEXT, passwd TEXT)''')
    threading.Thread(target=run_socket_server).start()
    app.run(port=5001)