276 lines
11 KiB
Python
276 lines
11 KiB
Python
import socket
|
|
import threading
|
|
import json
|
|
import sqlite3
|
|
from datetime import datetime
|
|
import os
|
|
|
|
class ChatServer:
|
|
def __init__(self, host='0.0.0.0', port=5555):
|
|
self.host = host
|
|
self.port = port
|
|
self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
self.server.bind((host, port))
|
|
self.server.listen()
|
|
self.clients = {}
|
|
self.groups = {'General': set()}
|
|
self.setup_database()
|
|
|
|
def setup_database(self):
|
|
# 确保数据库文件存在并初始化
|
|
db_exists = os.path.exists('chat_server.db')
|
|
self.db = sqlite3.connect('chat_server.db', check_same_thread=False)
|
|
cursor = self.db.cursor()
|
|
|
|
# 创建用户表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
username TEXT UNIQUE NOT NULL,
|
|
password TEXT NOT NULL
|
|
)
|
|
''')
|
|
|
|
# 创建消息表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
sender TEXT NOT NULL,
|
|
receiver TEXT NOT NULL,
|
|
message TEXT NOT NULL,
|
|
timestamp DATETIME NOT NULL,
|
|
is_group INTEGER NOT NULL
|
|
)
|
|
''')
|
|
|
|
# 如果是首次创建数据库，添加一个默认管理员用户
|
|
if not db_exists:
|
|
try:
|
|
cursor.execute('INSERT INTO users (username, password) VALUES (?, ?)',
|
|
('admin', 'admin123'))
|
|
print("Created default admin user: admin/admin123")
|
|
except sqlite3.IntegrityError:
|
|
pass
|
|
|
|
self.db.commit()
|
|
|
|
def authenticate_user(self, username, password):
|
|
cursor = self.db.cursor()
|
|
cursor.execute('SELECT * FROM users WHERE username=? AND password=?', (username, password))
|
|
return cursor.fetchone() is not None
|
|
|
|
def register_user(self, username, password):
|
|
if not username or not password:
|
|
return False
|
|
|
|
try:
|
|
cursor = self.db.cursor()
|
|
cursor.execute('INSERT INTO users (username, password) VALUES (?, ?)', (username, password))
|
|
self.db.commit()
|
|
return True
|
|
except sqlite3.IntegrityError:
|
|
return False
|
|
|
|
def save_message(self, sender, receiver, message, is_group):
|
|
cursor = self.db.cursor()
|
|
cursor.execute('''
|
|
INSERT INTO messages (sender, receiver, message, timestamp, is_group)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
''', (sender, receiver, message, datetime.now(), 1 if is_group else 0))
|
|
self.db.commit()
|
|
|
|
def send_user_list(self):
|
|
user_list = list(self.clients.keys())
|
|
for client in self.clients.values():
|
|
try:
|
|
client.send(json.dumps({
|
|
'type': 'user_list',
|
|
'users': user_list
|
|
}).encode('utf-8'))
|
|
except:
|
|
pass
|
|
|
|
def send_group_list(self, username):
|
|
group_list = [group for group, members in self.groups.items() if username in members]
|
|
if username in self.clients:
|
|
try:
|
|
self.clients[username].send(json.dumps({
|
|
'type': 'group_list',
|
|
'groups': group_list
|
|
}).encode('utf-8'))
|
|
except:
|
|
pass
|
|
|
|
def send_group_list_to_all(self):
|
|
for username in self.clients:
|
|
self.send_group_list(username)
|
|
|
|
def send_initial_messages(self, username):
|
|
cursor = self.db.cursor()
|
|
cursor.execute('''
|
|
SELECT sender, receiver, message, timestamp, is_group FROM messages
|
|
WHERE receiver=? OR sender=? OR (is_group=1 AND receiver IN (
|
|
SELECT name FROM sqlite_master WHERE type='table' AND name='groups'
|
|
))
|
|
ORDER BY timestamp
|
|
LIMIT 100
|
|
''', (username, username))
|
|
|
|
messages = []
|
|
for sender, receiver, message, timestamp, is_group in cursor.fetchall():
|
|
messages.append({
|
|
'sender': sender,
|
|
'receiver': receiver,
|
|
'message': message,
|
|
'timestamp': timestamp,
|
|
'is_group': is_group
|
|
})
|
|
|
|
if username in self.clients:
|
|
try:
|
|
self.clients[username].send(json.dumps({
|
|
'type': 'initial_messages',
|
|
'messages': messages
|
|
}).encode('utf-8'))
|
|
except:
|
|
pass
|
|
|
|
def user_exists(self, username):
|
|
cursor = self.db.cursor()
|
|
cursor.execute('SELECT 1 FROM users WHERE username=?', (username,))
|
|
return cursor.fetchone() is not None
|
|
|
|
def handle_client(self, client, address):
|
|
username = None
|
|
print(f"New connection from {address}")
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
message = client.recv(1024).decode('utf-8')
|
|
if not message:
|
|
break
|
|
|
|
data = json.loads(message)
|
|
print(f"Received from {address}: {data}")
|
|
|
|
if data['type'] == 'login':
|
|
username = data['username']
|
|
password = data['password']
|
|
if self.authenticate_user(username, password):
|
|
self.clients[username] = client
|
|
self.groups['General'].add(username)
|
|
client.send(json.dumps({'type': 'login_success'}).encode('utf-8'))
|
|
self.send_user_list()
|
|
self.send_group_list(username)
|
|
self.send_initial_messages(username)
|
|
print(f"User {username} logged in successfully")
|
|
else:
|
|
client.send(json.dumps({'type': 'login_fail', 'reason': 'Invalid username or password'}).encode('utf-8'))
|
|
print(f"Login failed for {username}")
|
|
|
|
elif data['type'] == 'register':
|
|
username = data['username']
|
|
password = data['password']
|
|
if len(username) < 3 or len(password) < 3:
|
|
client.send(json.dumps({'type': 'register_fail', 'reason': 'Username and password must be at least 3 characters'}).encode('utf-8'))
|
|
elif self.register_user(username, password):
|
|
client.send(json.dumps({'type': 'register_success'}).encode('utf-8'))
|
|
print(f"New user registered: {username}")
|
|
else:
|
|
client.send(json.dumps({'type': 'register_fail', 'reason': 'Username already taken'}).encode('utf-8'))
|
|
|
|
elif data['type'] == 'message':
|
|
sender = data['sender']
|
|
receiver = data['receiver']
|
|
msg = data['message']
|
|
is_group = data['is_group']
|
|
self.save_message(sender, receiver, msg, is_group)
|
|
if is_group:
|
|
for member in self.groups.get(receiver, set()):
|
|
if member in self.clients and member != sender:
|
|
try:
|
|
self.clients[member].send(json.dumps({
|
|
'type': 'message',
|
|
'sender': sender,
|
|
'receiver': receiver,
|
|
'message': msg,
|
|
'is_group': True
|
|
}).encode('utf-8'))
|
|
except:
|
|
pass
|
|
else:
|
|
if receiver in self.clients:
|
|
try:
|
|
self.clients[receiver].send(json.dumps({
|
|
'type': 'message',
|
|
'sender': sender,
|
|
'receiver': receiver,
|
|
'message': msg,
|
|
'is_group': False
|
|
}).encode('utf-8'))
|
|
except:
|
|
pass
|
|
|
|
elif data['type'] == 'create_group':
|
|
group_name = data['group_name']
|
|
if group_name not in self.groups:
|
|
self.groups[group_name] = set()
|
|
for user in data['members']:
|
|
if user in self.clients or self.user_exists(user):
|
|
self.groups[group_name].add(user)
|
|
self.send_group_list_to_all()
|
|
client.send(json.dumps({'type': 'group_created', 'group': group_name}).encode('utf-8'))
|
|
else:
|
|
client.send(json.dumps({'type': 'group_exists', 'group': group_name}).encode('utf-8'))
|
|
|
|
except json.JSONDecodeError:
|
|
print(f"Invalid JSON received from {address}")
|
|
break
|
|
except ConnectionResetError:
|
|
print(f"Connection reset by {address}")
|
|
break
|
|
except Exception as e:
|
|
print(f"Error handling client {address}: {str(e)}")
|
|
break
|
|
|
|
except Exception as e:
|
|
print(f"Client handler error for {address}: {str(e)}")
|
|
finally:
|
|
if username and username in self.clients:
|
|
print(f"User {username} disconnected")
|
|
del self.clients[username]
|
|
for group in self.groups.values():
|
|
if username in group:
|
|
group.remove(username)
|
|
self.send_user_list()
|
|
try:
|
|
client.close()
|
|
except:
|
|
pass
|
|
|
|
def start(self):
|
|
print(f"Server started on {self.host}:{self.port}")
|
|
print("Database initialized at chat_server.db")
|
|
try:
|
|
while True:
|
|
client, address = self.server.accept()
|
|
thread = threading.Thread(target=self.handle_client, args=(client, address))
|
|
thread.daemon = True
|
|
thread.start()
|
|
except KeyboardInterrupt:
|
|
print("\nServer is shutting down...")
|
|
finally:
|
|
for client in self.clients.values():
|
|
try:
|
|
client.close()
|
|
except:
|
|
pass
|
|
self.server.close()
|
|
self.db.close()
|
|
print("Server stopped")
|
|
|
|
if __name__ == "__main__":
|
|
server = ChatServer()
|
|
server.start() |