339 lines
No EOL
16 KiB
Python
339 lines
No EOL
16 KiB
Python
import socket
|
|
import select
|
|
import uuid
|
|
import pickle
|
|
import os
|
|
import threading
|
|
from flask import Flask, request, jsonify
|
|
from flask_socketio import SocketIO, emit, join_room, leave_room
|
|
|
|
# Flask app for REST API and WebSocket
|
|
app = Flask(__name__)
|
|
socketio = SocketIO(app, cors_allowed_origins="*") # Allow all origins
|
|
|
|
CHATPORT = 8181
|
|
BUFFER_SIZE = 1024
|
|
MAX_CONNECT_QUEUE = 10
|
|
DATA_FILE = 'server_data.pkl'
|
|
UUID_FILE = 'uuids.pkl'
|
|
|
|
# Global variables for chat server
|
|
socket_list = []
|
|
client_info = {}
|
|
chat_rooms = {}
|
|
client_states = {}
|
|
uuid_dict = {}
|
|
message_history = {} # New: Store message history for each chat room
|
|
server_socket = None
|
|
quitFlag = True
|
|
|
|
# Utility functions
|
|
def save_data(client_info, chat_rooms, client_states):
|
|
serializable_client_info = {info["user_id"]: info for sock, info in client_info.items()}
|
|
serializable_chat_rooms = {info["user_id"]: chat_rooms[sock] for sock, info in client_info.items()}
|
|
serializable_client_states = {info["user_id"]: client_states[sock] for sock, info in client_info.items()}
|
|
with open(DATA_FILE, 'wb') as f:
|
|
pickle.dump((serializable_client_info, serializable_chat_rooms, serializable_client_states), f)
|
|
|
|
def load_data():
|
|
if os.path.exists(DATA_FILE):
|
|
with open(DATA_FILE, 'rb') as f:
|
|
serializable_client_info, serializable_chat_rooms, serializable_client_states = pickle.load(f)
|
|
client_info = {uuid.UUID(user_id): info for user_id, info in serializable_client_info.items()}
|
|
chat_rooms = {uuid.UUID(user_id): room for user_id, room in serializable_chat_rooms.items()}
|
|
client_states = {uuid.UUID(user_id): state for user_id, state in serializable_client_states.items()}
|
|
return client_info, chat_rooms, client_states
|
|
return {}, {}, {}
|
|
|
|
def save_uuids(uuid_dict):
|
|
with open(UUID_FILE, 'wb') as f:
|
|
pickle.dump(uuid_dict, f)
|
|
|
|
def load_uuids():
|
|
if os.path.exists(UUID_FILE):
|
|
with open(UUID_FILE, 'rb') as f:
|
|
return pickle.load(f)
|
|
return {}
|
|
|
|
def get_client_uuid(client_socket, uuid_dict):
|
|
client_address = client_socket.getpeername()
|
|
client_key = f"{client_address[0]}_{client_address[1]}"
|
|
|
|
if client_key in uuid_dict:
|
|
user_id = uuid_dict[client_key]
|
|
else:
|
|
user_id = str(uuid.uuid4())
|
|
uuid_dict[client_key] = user_id
|
|
save_uuids(uuid_dict)
|
|
|
|
return user_id
|
|
|
|
# REST API Endpoints
|
|
@app.route('/send_message', methods=['POST'])
|
|
def send_message():
|
|
"""Send a message to a specific chat room."""
|
|
data = request.json
|
|
username = data.get('username')
|
|
chat_room = data.get('chat_room', 'default')
|
|
message = data.get('message')
|
|
|
|
if not username or not message:
|
|
return jsonify({"error": "Username and message are required"}), 400
|
|
|
|
# Store the message in the chat room's history
|
|
if chat_room not in message_history:
|
|
message_history[chat_room] = []
|
|
message_history[chat_room].append(f"{username}: {message}")
|
|
|
|
# Broadcast the message to all clients in the same chat room
|
|
send_message = f"{username} >> {message}".encode('utf-8')
|
|
for client_sock in socket_list:
|
|
if client_sock != server_socket and chat_rooms.get(client_sock) == chat_room:
|
|
client_sock.send(send_message)
|
|
|
|
return jsonify({"status": "Message sent and saved in history"}), 200
|
|
|
|
@app.route('/create_chatroom', methods=['POST'])
|
|
def create_chatroom():
|
|
"""Create a new chat room."""
|
|
data = request.json
|
|
chat_room = data.get('chat_room')
|
|
|
|
if not chat_room:
|
|
return jsonify({"error": "Chat room name is required"}), 400
|
|
|
|
# Add logic to create a chat room (if needed)
|
|
return jsonify({"status": f"Chat room '{chat_room}' created"}), 200
|
|
|
|
@app.route('/get_messages', methods=['GET'])
|
|
def get_messages():
|
|
"""Retrieve messages for a specific user (not implemented here)."""
|
|
# You can implement a message queue or buffer for each user to store messages
|
|
return jsonify({"error": "Not implemented"}), 501
|
|
|
|
@app.route('/get_history', methods=['GET'])
|
|
def get_history():
|
|
"""Retrieve message history for a specific chat room."""
|
|
chat_room = request.args.get('chat_room', 'default')
|
|
if chat_room not in message_history:
|
|
return jsonify({"error": f"No history found for chat room '{chat_room}'"}), 404
|
|
|
|
return jsonify({"chat_room": chat_room, "messages": message_history[chat_room]}), 200
|
|
|
|
# WebSocket event handlers
|
|
@socketio.on('connect')
|
|
def handle_connect():
|
|
print("WebSocket client connected")
|
|
|
|
@socketio.on('disconnect')
|
|
def handle_disconnect():
|
|
print("WebSocket client disconnected")
|
|
|
|
@socketio.on('join')
|
|
def handle_join(data):
|
|
print(f"Received join event: {data}")
|
|
username = data.get('username')
|
|
chat_room = data.get('chat_room', 'default')
|
|
|
|
if not username:
|
|
emit('error', {'message': 'Username is required'})
|
|
return
|
|
|
|
join_room(chat_room)
|
|
emit('message', {'username': 'Server', 'message': f"{username} has joined the chat room '{chat_room}'."}, room=chat_room)
|
|
|
|
@socketio.on('leave')
|
|
def handle_leave(data):
|
|
username = data.get('username')
|
|
chat_room = data.get('chat_room', 'default')
|
|
leave_room(chat_room)
|
|
emit('message', {'message': f"{username} has left the chat room '{chat_room}'."}, room=chat_room)
|
|
|
|
@socketio.on('send_message')
|
|
def handle_send_message(data):
|
|
print(f"Received send_message event: {data}")
|
|
username = data.get('username')
|
|
chat_room = data.get('chat_room', 'default')
|
|
message = data.get('message')
|
|
|
|
if not username or not message:
|
|
emit('error', {'message': 'Username and message are required'})
|
|
return
|
|
|
|
# Store the message in the chat room's history
|
|
if chat_room not in message_history:
|
|
message_history[chat_room] = []
|
|
message_history[chat_room].append(f"{username}: {message}")
|
|
|
|
# Broadcast the message to the chat room
|
|
emit('message', {'username': username, 'message': message}, room=chat_room)
|
|
|
|
# Chat server logic
|
|
def chat_server():
|
|
global quitFlag, server_socket, socket_list, client_info, chat_rooms, client_states, uuid_dict
|
|
|
|
# Create and configure the server socket
|
|
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
|
|
server_socket.bind(('', CHATPORT))
|
|
server_socket.listen(MAX_CONNECT_QUEUE)
|
|
server_socket.setblocking(False)
|
|
|
|
socket_list.append(server_socket)
|
|
print(f"Chat server listening on port {CHATPORT}")
|
|
|
|
while quitFlag:
|
|
read_sockets, _, _ = select.select(socket_list, [], [])
|
|
for sock in read_sockets:
|
|
if sock == server_socket:
|
|
# Accept new client connections
|
|
client_socket, client_address = server_socket.accept()
|
|
client_socket.setblocking(False)
|
|
socket_list.append(client_socket)
|
|
print(f"Accepted connection from {client_address}")
|
|
client_socket.send(b"Welcome to the chat server!\n")
|
|
client_socket.send(b"Type in a username: \n")
|
|
client_states[client_socket] = "WAITING_FOR_USERNAME"
|
|
user_id = get_client_uuid(client_socket, uuid_dict)
|
|
client_info[client_socket] = {"user_id": user_id}
|
|
else:
|
|
# Handle client messages
|
|
try:
|
|
message = sock.recv(BUFFER_SIZE).decode('utf-8')
|
|
if message:
|
|
state = client_states.get(sock, "NORMAL")
|
|
if state == "WAITING_FOR_USERNAME":
|
|
username = message.strip()
|
|
if username == "admin":
|
|
sock.send(b"\nEnter admin password: \n")
|
|
client_states[sock] = "WAITING_FOR_ADMIN_PASSWORD"
|
|
client_info[sock]["username"] = username
|
|
else:
|
|
client_info[sock]["username"] = username
|
|
chat_rooms[sock] = "default"
|
|
sock.send(b"\nUsername accepted. You are now in the 'default' chatroom.\n")
|
|
client_states[sock] = "NORMAL"
|
|
elif state == "WAITING_FOR_ADMIN_PASSWORD":
|
|
password = message.strip()
|
|
if password == "adminpass":
|
|
sock.send(b"\nAdmin authenticated successfully.\n")
|
|
client_info[sock]["is_admin"] = True
|
|
chat_rooms[sock] = "default"
|
|
client_states[sock] = "NORMAL"
|
|
else:
|
|
sock.send(b"\nInvalid admin password. Please enter a username: \n")
|
|
client_states[sock] = "WAITING_FOR_USERNAME"
|
|
elif state == "WAITING_FOR_NEW_USERNAME":
|
|
client_info[sock]["username"] = message.strip()
|
|
sock.send(b"\nUsername changed successfully.\n")
|
|
print(f"\nClient switched username to {client_info[sock]['username']}.")
|
|
client_states[sock] = "NORMAL"
|
|
elif state == "WAITING_FOR_NEW_CHATROOM":
|
|
old_chatroom = chat_rooms[sock]
|
|
chat_rooms[sock] = message.strip()
|
|
sock.send(f"\nSwitched to chatroom '{chat_rooms[sock]}.'\n".encode('utf-8'))
|
|
print(f"\nClient switched to chatroom '{chat_rooms[sock]}.'")
|
|
if old_chatroom != "default" and old_chatroom not in chat_rooms.values():
|
|
print(f"\nChatroom '{old_chatroom}' is empty and removed.")
|
|
client_states[sock] = "NORMAL"
|
|
else:
|
|
if message.strip() == "%!":
|
|
sock.close()
|
|
socket_list.remove(sock)
|
|
del client_info[sock]
|
|
del chat_rooms[sock]
|
|
del client_states[sock]
|
|
print("\nClient has closed connection.")
|
|
elif message.strip() == "%?":
|
|
print("\nClient requested help.")
|
|
sock.send(b"Help: (%?)\n")
|
|
sock.send(b"Send messages like normal. Use enter to send a message.\n")
|
|
sock.send(b"Available Commands:\n")
|
|
sock.send(b"\t%? : Show this help prompt.\n")
|
|
sock.send(b"\t%! : Exit from the Chat Server.\n")
|
|
sock.send(b"\t%* : Switch to another chatroom.\n")
|
|
sock.send(b"\t%$ : Change Username.\n")
|
|
sock.send(b"\t%@ <username> <message> : Send a direct message to a user.\n")
|
|
if client_info[sock].get("is_admin"):
|
|
sock.send(b"\t%stop : Stop the server.\n")
|
|
elif message.strip() == "%*":
|
|
sock.send(b"\nEnter chatroom name: \n")
|
|
client_states[sock] = "WAITING_FOR_NEW_CHATROOM"
|
|
elif message.strip() == "%$":
|
|
sock.send(b"\nEnter new username: \n")
|
|
client_states[sock] = "WAITING_FOR_NEW_USERNAME"
|
|
elif message.strip().startswith("%@ "):
|
|
parts = message.strip().split(' ', 2)
|
|
if len(parts) < 3:
|
|
sock.send(b"\nInvalid direct message format. Use: %@ <username> <message>\n")
|
|
else:
|
|
target_username = parts[1]
|
|
dm_message = parts[2]
|
|
target_sock = None
|
|
for s, info in client_info.items():
|
|
if info.get("username") == target_username:
|
|
target_sock = s
|
|
break
|
|
if target_sock:
|
|
sender_username = client_info[sock]["username"]
|
|
target_sock.send(f"DM from {sender_username}: {dm_message}\n".encode('utf-8'))
|
|
sock.send(b"\nDirect message sent.\n")
|
|
else:
|
|
sock.send(b"\nUser not found.\n")
|
|
elif message.strip() == "%stop" and client_info[sock].get("is_admin"):
|
|
print("\nAdmin has stopped the server.")
|
|
quitFlag = False
|
|
else:
|
|
username = client_info[sock]["username"]
|
|
chat_room = chat_rooms[sock]
|
|
send_message = f"{username} >> {message}".encode('utf-8')
|
|
|
|
# Store the message in the chat room's history
|
|
if chat_room not in message_history:
|
|
message_history[chat_room] = []
|
|
message_history[chat_room].append(f"{username}: {message}")
|
|
|
|
# Broadcast the message to all clients in the same chat room
|
|
for client_sock in socket_list:
|
|
if client_sock != server_socket and chat_rooms.get(client_sock) == chat_room:
|
|
client_sock.send(send_message)
|
|
else:
|
|
sock.close()
|
|
socket_list.remove(sock)
|
|
if sock in client_info:
|
|
del client_info[sock]
|
|
if sock in chat_rooms:
|
|
old_chatroom = chat_rooms[sock]
|
|
del chat_rooms[sock]
|
|
if old_chatroom != "default" and old_chatroom not in chat_rooms.values():
|
|
print(f"\nChatroom '{old_chatroom}' is empty and has been removed.")
|
|
if sock in client_states:
|
|
del client_states[sock]
|
|
print("\nClient has closed connection.")
|
|
except Exception as err:
|
|
print(f"Error: {err}")
|
|
sock.close()
|
|
socket_list.remove(sock)
|
|
if sock in client_info:
|
|
del client_info[sock]
|
|
if sock in chat_rooms:
|
|
old_chatroom = chat_rooms[sock]
|
|
del chat_rooms[sock]
|
|
if old_chatroom != "default" and old_chatroom not in chat_rooms.values():
|
|
print(f"\nChatroom '{old_chatroom}' is empty and has been removed.")
|
|
if sock in client_states:
|
|
del client_states[sock]
|
|
print("\nClient has closed connection.")
|
|
|
|
# Save data before shutting down
|
|
save_data(client_info, chat_rooms, client_states)
|
|
server_socket.close()
|
|
|
|
if __name__ == "__main__":
|
|
# Start the chat server in a separate thread
|
|
server_thread = threading.Thread(target=chat_server, daemon=True)
|
|
server_thread.start()
|
|
|
|
# Start the Flask REST API and WebSocket server
|
|
socketio.run(app, host='0.0.0.0', port=5000) |