474 lines
No EOL
20 KiB
Python
474 lines
No EOL
20 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
|
|
from flask_cors import CORS
|
|
from datetime import datetime # Import for timestamps
|
|
import markdown # Import the markdown library
|
|
|
|
# Flask app for REST API and WebSocket
|
|
app = Flask(__name__)
|
|
socketio = SocketIO(app, cors_allowed_origins="*") # Allow all origins
|
|
|
|
# Enable CORS for the Flask app
|
|
CORS(app)
|
|
|
|
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
|
|
|
|
def handle_message(username, chat_room, message):
|
|
"""Store the message in history and broadcast it to all clients."""
|
|
# Add a timestamp to the message
|
|
timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
|
|
|
# Convert the message to HTML using Markdown
|
|
html_message = markdown.markdown(message)
|
|
|
|
# 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"[{timestamp}] {username}: {message}")
|
|
|
|
# Broadcast the message to WebSocket clients in the chat room
|
|
socketio.emit('message', {
|
|
'username': username,
|
|
'message': html_message, # Send the HTML-rendered message
|
|
'timestamp': timestamp,
|
|
'chat_room': chat_room
|
|
}, room=chat_room)
|
|
|
|
# Broadcast the message to socket-based clients in the chat room
|
|
send_message = f"[{timestamp}] {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)
|
|
|
|
def handle_command(username, chat_room, command, sock=None):
|
|
"""Process commands sent by clients."""
|
|
if command.strip() == "%?":
|
|
help_message = (
|
|
"Help: (%?)\n"
|
|
"Send messages like normal. Use enter to send a message.\n"
|
|
"Available Commands:\n"
|
|
"\t%? : Show this help prompt.\n"
|
|
"\t%! : Exit from the Chat Server.\n"
|
|
"\t%* : Switch to another chatroom.\n"
|
|
"\t%$ : Change Username.\n"
|
|
"\t%@ <username> <message> : Send a direct message to a user.\n"
|
|
)
|
|
if sock:
|
|
sock.send(help_message.encode('utf-8'))
|
|
else:
|
|
socketio.emit('message', {'username': 'Server', 'message': help_message}, room=chat_room)
|
|
return
|
|
|
|
if command.strip() == "%*":
|
|
if sock:
|
|
sock.send(b"\nEnter chatroom name: \n")
|
|
client_states[sock] = "WAITING_FOR_NEW_CHATROOM"
|
|
else:
|
|
socketio.emit('message', {'username': 'Server', 'message': "Enter chatroom name:"}, room=chat_room)
|
|
return
|
|
|
|
if command.strip() == "%$":
|
|
if sock:
|
|
sock.send(b"\nEnter new username: \n")
|
|
client_states[sock] = "WAITING_FOR_NEW_USERNAME"
|
|
else:
|
|
socketio.emit('message', {'username': 'Server', 'message': "Enter new username:"}, room=chat_room)
|
|
return
|
|
|
|
if command.strip().startswith("%@ "):
|
|
parts = command.strip().split(' ', 2)
|
|
if len(parts) < 3:
|
|
error_message = "Invalid direct message format. Use: %@ <username> <message>"
|
|
if sock:
|
|
sock.send(f"\n{error_message}\n".encode('utf-8'))
|
|
else:
|
|
socketio.emit('message', {'username': 'Server', 'message': error_message}, room=chat_room)
|
|
else:
|
|
target_username = parts[1]
|
|
dm_message = parts[2]
|
|
timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
|
|
|
# Check if the target is a socket-based client
|
|
target_sock = None
|
|
for s, info in client_info.items():
|
|
if not info.get("is_websocket") and info.get("username") == target_username:
|
|
target_sock = s
|
|
break
|
|
|
|
# Check if the target is a WebSocket client
|
|
target_websocket_sid = None
|
|
for sid, info in client_info.items():
|
|
if info.get("is_websocket") and info.get("username") == target_username:
|
|
target_websocket_sid = sid
|
|
break
|
|
|
|
if target_sock:
|
|
# Send the direct message to the socket-based client
|
|
sender_username = username
|
|
target_sock.send(f"[{timestamp}] DM from {sender_username}: {dm_message}\n".encode('utf-8'))
|
|
if sock:
|
|
sock.send(f"[{timestamp}] Direct message sent.\n".encode('utf-8'))
|
|
elif target_websocket_sid:
|
|
# Send the direct message to the WebSocket client
|
|
socketio.emit('message', {
|
|
'username': username,
|
|
'message': dm_message,
|
|
'timestamp': timestamp,
|
|
'chat_room': 'Direct Message'
|
|
}, to=target_websocket_sid)
|
|
if sock:
|
|
sock.send(f"[{timestamp}] Direct message sent.\n".encode('utf-8'))
|
|
else:
|
|
# Target user not found
|
|
error_message = "User not found."
|
|
if sock:
|
|
sock.send(f"\n{error_message}\n".encode('utf-8'))
|
|
else:
|
|
socketio.emit('message', {'username': 'Server', 'message': error_message}, room=chat_room)
|
|
return
|
|
|
|
if command.strip() == "%stop":
|
|
if client_info.get(sock, {}).get("is_admin"):
|
|
print("\nAdmin has stopped the server.")
|
|
global quitFlag
|
|
quitFlag = False
|
|
if sock:
|
|
sock.send(b"\nServer is stopping...\n")
|
|
else:
|
|
socketio.emit('message', {'username': 'Server', 'message': "Server is stopping..."}, room=chat_room)
|
|
else:
|
|
error_message = "You do not have permission to stop the server."
|
|
if sock:
|
|
sock.send(f"\n{error_message}\n".encode('utf-8'))
|
|
else:
|
|
socketio.emit('message', {'username': 'Server', 'message': error_message}, room=chat_room)
|
|
return
|
|
|
|
# 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
|
|
|
|
if message.strip().startswith("%"):
|
|
handle_command(username, chat_room, message)
|
|
else:
|
|
handle_message(username, chat_room, message)
|
|
|
|
return jsonify({"status": "Message sent and processed"}), 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({"messages": []}) # Return an empty list if no history exists
|
|
|
|
# Return the full message history with timestamps
|
|
return jsonify({"messages": message_history[chat_room]})
|
|
|
|
@app.route('/send_direct_message', methods=['POST'])
|
|
def send_direct_message():
|
|
"""Send a direct message to a specific user."""
|
|
data = request.json
|
|
sender_username = data.get('sender_username')
|
|
target_username = data.get('target_username')
|
|
message = data.get('message')
|
|
|
|
if not sender_username or not target_username or not message:
|
|
return jsonify({"error": "Sender username, target username, and message are required"}), 400
|
|
|
|
timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
|
|
|
# Check if the target is a socket-based client
|
|
target_sock = None
|
|
for s, info in client_info.items():
|
|
if not info.get("is_websocket") and info.get("username") == target_username:
|
|
target_sock = s
|
|
break
|
|
|
|
# Check if the target is a WebSocket client
|
|
target_websocket_sid = None
|
|
for sid, info in client_info.items():
|
|
if info.get("is_websocket") and info.get("username") == target_username:
|
|
target_websocket_sid = sid
|
|
break
|
|
|
|
if target_sock:
|
|
# Send the direct message to the socket-based client
|
|
target_sock.send(f"[{timestamp}] DM from {sender_username}: {message}\n".encode('utf-8'))
|
|
return jsonify({"status": "Direct message sent to socket-based client"}), 200
|
|
elif target_websocket_sid:
|
|
# Send the direct message to the WebSocket client
|
|
socketio.emit('message', {
|
|
'username': sender_username,
|
|
'message': message,
|
|
'timestamp': timestamp,
|
|
'chat_room': 'Direct Message'
|
|
}, to=target_websocket_sid)
|
|
return jsonify({"status": "Direct message sent to WebSocket client"}), 200
|
|
else:
|
|
# Target user not found
|
|
return jsonify({"error": "Target user not found"}), 404
|
|
|
|
# WebSocket event handlers
|
|
@socketio.on('connect')
|
|
def handle_connect():
|
|
print("WebSocket client connected")
|
|
# Assign a unique ID to the WebSocket client
|
|
user_id = str(uuid.uuid4())
|
|
client_info[request.sid] = {"user_id": user_id, "username": None, "is_websocket": True}
|
|
print(f"WebSocket client registered with ID: {user_id}")
|
|
|
|
@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
|
|
|
|
# Update client info for WebSocket clients
|
|
client_info[request.sid]["username"] = username
|
|
client_info[request.sid]["chat_room"] = chat_room
|
|
|
|
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
|
|
|
|
if message.strip().startswith("%"):
|
|
handle_command(username, chat_room, message)
|
|
else:
|
|
handle_message(username, chat_room, message)
|
|
|
|
# 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.")
|
|
else:
|
|
if message.strip().startswith("%"):
|
|
handle_command(client_info[sock]["username"], chat_rooms[sock], message, sock)
|
|
else:
|
|
username = client_info[sock]["username"]
|
|
chat_room = chat_rooms[sock]
|
|
handle_message(username, chat_room, 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) |