加入token和HTTP API聊天功能

CS3.1.py

@@ -3,7 +3,9 @@ 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()
@@ -11,6 +13,7 @@ socket_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 
 active_users = {}
 chat_connections = []
+tokens = {}
 
 def get_db_connection():
     conn = sqlite3.connect("usr.db")
@@ -60,6 +63,18 @@ def register_user(usr, pwd):
         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()
@@ -71,6 +86,30 @@ def register1():
     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:
@@ -92,8 +131,9 @@ def handle_socket_message(data, addr, conn):
             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"}
+                    return {"type": "login_1", "status": "success", "message": "Login successful", "token": generate_token(data['username'])}
             return {"type": "login_0", "status": "error", "message": "Invalid credentials"}
         elif action == 'chat':
             if addr[0] in active_users:
@@ -103,7 +143,7 @@ def handle_socket_message(data, addr, conn):
                     "message": data['message']
                 }
                 broadcast_message(message)
-                return {"type": "chat", "status": "success", }
+                return {"type": "chat", "status": "success"}
             return {"type": "chat", "status": "error", "message": "Not logged in"}
     except Exception as e:
         return {"status": "error", "message": str(e)}