back-end / server.py
zeegeedee's picture
Update server.py
93b40b2 verified
Raw
History Blame Contribute Delete
24.4 kB
import random
from funcs import *
the_reachable_void()
import asyncio
import json
import os
import uuid
import bcrypt
import time
import secrets
import hashlib
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import Response
from fastapi.middleware.cors import CORSMiddleware
import urllib.request
try:
from huggingface_hub import HfApi, hf_hub_url, hf_hub_download
except ImportError:
HfApi = None
hf_hub_url = None
hf_hub_download = None
# Environment variables configuration for Hugging Face storage
HF_TOKEN = os.getenv("database_token")
DATASET_REPO_ID = os.getenv("dataset_link")
USERS_FILENAME = "users.json"
CONVERSATIONS_FILENAME = "conversations.json"
SALTS_FILENAME = "remember_me_salts.json"
RECENTS_FILENAME = "users_recent_chats.json"
OFFLINE_QUEUES_FILENAME = "offline_queues.json"
os.makedirs("/data", exist_ok=True)
LOCAL_USERS_PATH = os.path.join("/data", USERS_FILENAME)
LOCAL_CONVERSATIONS_PATH = os.path.join("/data", CONVERSATIONS_FILENAME)
LOCAL_SALTS_PATH = os.path.join("/data", SALTS_FILENAME)
LOCAL_RECENTS_PATH = os.path.join("/data", RECENTS_FILENAME)
LOCAL_OFFLINE_QUEUES_PATH = os.path.join("/data", OFFLINE_QUEUES_FILENAME)
hf_api = HfApi(token=HF_TOKEN) if HF_TOKEN and HfApi else None
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
MAINTENANCE_FILENAME = "maintenance.json"
LOCAL_MAINTENANCE_PATH = os.path.join("/data", MAINTENANCE_FILENAME)
PENDING_USERS_SYNC = False
PENDING_CONVERSATIONS_SYNC = False
PENDING_SALTS_SYNC = False
PENDING_RECENTS_SYNC = False
PENDING_OFFLINE_QUEUES_SYNC = False
def load_maintenance():
if os.path.exists(LOCAL_MAINTENANCE_PATH):
with open(LOCAL_MAINTENANCE_PATH, "r", encoding="utf-8") as f:
return json.load(f).get("maintenance", False)
return False
@app.get("/maintenance")
async def get_maintenance():
return {"maintenance": load_maintenance()}
@app.get("/")
@app.head("/")
async def health_check():
return Response(content="OK", media_type="text/plain")
def hash_password(plain_text):
salt = bcrypt.gensalt()
return bcrypt.hashpw(plain_text.encode('utf-8'), salt).decode('utf-8')
def verify_password(plain_text, hashed_text):
try:
return bcrypt.checkpw(plain_text.encode('utf-8'), hashed_text.encode('utf-8'))
except Exception:
return False
file_lock = asyncio.Lock()
def conversation_key(user_a, user_b):
return "|".join(sorted([user_a, user_b]))
def fetch_hf_file(filename):
if not hf_hub_url or not HF_TOKEN or not DATASET_REPO_ID:
raise RuntimeError("Hugging Face storage is not configured")
url = hf_hub_url(repo_id=DATASET_REPO_ID, filename=filename, repo_type="dataset")
req = urllib.request.Request(url, headers={"Authorization": f"Bearer {HF_TOKEN}"})
with urllib.request.urlopen(req) as response:
return response.read().decode("utf-8")
def load_databases():
global USERS_DB, CONVERSATIONS_DB, SALTS_DB, RECENTS_DB, OFFLINE_QUEUES_DB, NEXT_USER_ID
# Load offline queues safely
if os.path.exists(LOCAL_OFFLINE_QUEUES_PATH):
try:
with open(LOCAL_OFFLINE_QUEUES_PATH, "r", encoding="utf-8") as f:
OFFLINE_QUEUES_DB = json.load(f)
except Exception:
OFFLINE_QUEUES_DB = {}
elif HF_TOKEN and DATASET_REPO_ID and hf_hub_download:
try:
downloaded_path = hf_hub_download(repo_id=DATASET_REPO_ID, filename=OFFLINE_QUEUES_FILENAME, repo_type="dataset", token=HF_TOKEN)
with open(downloaded_path, "r", encoding="utf-8") as f:
OFFLINE_QUEUES_DB = json.load(f)
with open(LOCAL_OFFLINE_QUEUES_PATH, "w", encoding="utf-8") as f:
json.dump(OFFLINE_QUEUES_DB, f, indent=4)
except Exception:
OFFLINE_QUEUES_DB = {}
else:
OFFLINE_QUEUES_DB = {}
if not HF_TOKEN or not DATASET_REPO_ID:
print("Running in local-only mode.")
if os.path.exists(LOCAL_USERS_PATH):
with open(LOCAL_USERS_PATH, "r", encoding="utf-8") as f:
raw = json.load(f)
USERS_DB = raw.get("users", raw)
else:
USERS_DB = {}
if os.path.exists(LOCAL_CONVERSATIONS_PATH):
with open(LOCAL_CONVERSATIONS_PATH, "r", encoding="utf-8") as f:
CONVERSATIONS_DB = json.load(f)
else:
CONVERSATIONS_DB = {}
if os.path.exists(LOCAL_SALTS_PATH):
with open(LOCAL_SALTS_PATH, "r", encoding="utf-8") as f:
SALTS_DB = json.load(f)
else:
SALTS_DB = {}
if os.path.exists(LOCAL_RECENTS_PATH):
with open(LOCAL_RECENTS_PATH, "r", encoding="utf-8") as f:
RECENTS_DB = json.load(f)
else:
RECENTS_DB = {}
existing_ids = [int(u["user_id"]) for u in USERS_DB.values() if str(u.get("user_id", "")).isdigit()]
NEXT_USER_ID = max(existing_ids, default=0) + 1
return
try:
raw_users = fetch_hf_file(USERS_FILENAME)
parsed = json.loads(raw_users)
USERS_DB = parsed.get("users", parsed)
with open(LOCAL_USERS_PATH, "w", encoding="utf-8") as f:
json.dump(USERS_DB, f, indent=4)
print(f"Loaded {len(USERS_DB)} user(s) from HF.")
except Exception as e:
print(f"Failed to load users from HF: {e}")
USERS_DB = {}
try:
raw_convos = fetch_hf_file(CONVERSATIONS_FILENAME)
CONVERSATIONS_DB = json.loads(raw_convos)
with open(LOCAL_CONVERSATIONS_PATH, "w", encoding="utf-8") as f:
json.dump(CONVERSATIONS_DB, f, indent=4)
print(f"Loaded {len(CONVERSATIONS_DB)} conversation(s) from HF.")
except Exception as e:
print(f"No conversations file found or failed to load: {e}")
CONVERSATIONS_DB = {}
try:
raw_salts = fetch_hf_file(SALTS_FILENAME)
SALTS_DB = json.loads(raw_salts)
with open(LOCAL_SALTS_PATH, "w", encoding="utf-8") as f:
json.dump(SALTS_DB, f, indent=4)
print(f"Loaded {len(SALTS_DB)} salt(s) from HF.")
except Exception as e:
print(f"No salts file found or failed to load: {e}")
SALTS_DB = {}
try:
raw_recents = fetch_hf_file(RECENTS_FILENAME)
RECENTS_DB = json.loads(raw_recents)
with open(LOCAL_RECENTS_PATH, "w", encoding="utf-8") as f:
json.dump(RECENTS_DB, f, indent=4)
print(f"loaded {len(RECENTS_DB)} recent(s) from HF.")
except Exception as e:
print(f"No recents file or faile to load: {e}")
RECENTS_DB = {}
existing_ids = [int(u["user_id"]) for u in USERS_DB.values() if str(u.get("user_id", "")).isdigit()]
NEXT_USER_ID = max(existing_ids, default=0) + 1
load_databases()
MAX_RECENT_CONVERSATIONS_PER_USER = 100
def update_recent_chat_entry(user, partner, message_text, timestamp, display_name):
if not user or not partner:
return
if user not in RECENTS_DB:
RECENTS_DB[user] = []
entry = {
"partner": partner,
"last_message": message_text,
"timestamp": timestamp,
"display_name": display_name or partner,
}
RECENTS_DB[user] = [item for item in RECENTS_DB[user] if item.get("partner") != partner]
RECENTS_DB[user].append(entry)
RECENTS_DB[user].sort(key=lambda item: item.get("timestamp", 0), reverse=True)
RECENTS_DB[user] = RECENTS_DB[user][:MAX_RECENT_CONVERSATIONS_PER_USER]
def save_users_local():
global PENDING_USERS_SYNC
with open(LOCAL_USERS_PATH, "w", encoding="utf-8") as f:
json.dump(USERS_DB, f, indent=4)
PENDING_USERS_SYNC = True
def save_conversations_local():
global PENDING_CONVERSATIONS_SYNC
with open(LOCAL_CONVERSATIONS_PATH, "w", encoding="utf-8") as f:
json.dump(CONVERSATIONS_DB, f, indent=4)
PENDING_CONVERSATIONS_SYNC = True
def save_salts_local():
global PENDING_SALTS_SYNC
with open(LOCAL_SALTS_PATH, "w", encoding="utf-8") as f:
json.dump(SALTS_DB, f, indent=4)
PENDING_SALTS_SYNC = True
def save_recents_local():
global PENDING_RECENTS_SYNC
with open(LOCAL_RECENTS_PATH, "w", encoding="utf-8") as f:
json.dump(RECENTS_DB, f, indent=4)
PENDING_RECENTS_SYNC = True
def save_offline_queues_local():
global PENDING_OFFLINE_QUEUES_SYNC
with open(LOCAL_OFFLINE_QUEUES_PATH, "w", encoding="utf-8") as f:
json.dump(OFFLINE_QUEUES_DB, f, indent=4)
PENDING_OFFLINE_QUEUES_SYNC = True
async def background_cloud_sync():
global PENDING_USERS_SYNC, PENDING_CONVERSATIONS_SYNC, PENDING_SALTS_SYNC, PENDING_RECENTS_SYNC, PENDING_OFFLINE_QUEUES_SYNC
while True:
await asyncio.sleep(30)
if hf_api and DATASET_REPO_ID:
try:
if PENDING_USERS_SYNC:
await asyncio.to_thread(hf_api.upload_file, path_or_fileobj=LOCAL_USERS_PATH, path_in_repo=USERS_FILENAME, repo_id=DATASET_REPO_ID, repo_type="dataset")
PENDING_USERS_SYNC = False
if PENDING_CONVERSATIONS_SYNC:
await asyncio.to_thread(hf_api.upload_file, path_or_fileobj=LOCAL_CONVERSATIONS_PATH, path_in_repo=CONVERSATIONS_FILENAME, repo_id=DATASET_REPO_ID, repo_type="dataset")
PENDING_CONVERSATIONS_SYNC = False
if PENDING_SALTS_SYNC:
await asyncio.to_thread(hf_api.upload_file, path_or_fileobj=LOCAL_SALTS_PATH, path_in_repo=SALTS_FILENAME, repo_id=DATASET_REPO_ID, repo_type="dataset")
PENDING_SALTS_SYNC = False
if PENDING_RECENTS_SYNC:
await asyncio.to_thread(hf_api.upload_file, path_or_fileobj=LOCAL_RECENTS_PATH, path_in_repo=RECENTS_FILENAME, repo_id=DATASET_REPO_ID, repo_type="dataset")
PENDING_RECENTS_SYNC = False
if PENDING_OFFLINE_QUEUES_SYNC:
await asyncio.to_thread(hf_api.upload_file, path_or_fileobj=LOCAL_OFFLINE_QUEUES_PATH, path_in_repo=OFFLINE_QUEUES_FILENAME, repo_id=DATASET_REPO_ID, repo_type="dataset")
PENDING_OFFLINE_QUEUES_SYNC = False
except Exception as e:
print(f"[-] Background cloud sync partial failure: {e}")
@app.on_event("startup")
async def startup_event():
asyncio.create_task(background_cloud_sync())
ONLINE_USERS = {}
ACTIVE_TOKENS = {}
@app.websocket("/ws")
async def chat_handler(websocket: WebSocket):
global NEXT_USER_ID
await websocket.accept()
username = None
is_authenticated = False
try:
while True:
raw_data = await websocket.receive_text()
try:
auth_data = json.loads(raw_data)
except json.JSONDecodeError:
await websocket.send_text("[-] ERROR: Invalid JSON format.")
continue
action = auth_data.get("action")
username_input_raw = auth_data.get("user").lower() if auth_data.get("user") else ""
username_input_cleaned = username_input_raw.lower()
password = auth_data.get("pass")
display_name = auth_data.get("display", username_input_cleaned)
if action == "signup":
if not username_input_cleaned or not password:
await websocket.send_text("[-] ERROR: Username or password cannot be blank.")
continue
if username_input_cleaned in USERS_DB:
await websocket.send_text("[-] FAIL: Username already taken.")
continue
if len(password) >= 8:
banned_keywords = ["admin", "websitecreator", "thewebsitecreator", "webs1tecreator", "thewebs1tecreator"]
banned_chars = ["{", "}", "(", ")", "[", "]", "@", "#", ";", ":", "<", ">", "&", '"', "'", "\\"]
cleaned_display_name = display_name.lower().replace(" ", "")
cleaned_username = username_input_cleaned.lower().replace(" ", "")
if len(display_name) > 30:
await websocket.send_text("[-] ERROR: Display name is too long.")
continue
elif len(username_input_cleaned) > 30:
await websocket.send_text("[-] ERROR: username is too long.")
continue
elif any(char in display_name for char in banned_chars):
await websocket.send_text(r"""[-] ERROR: Display name contains banned character/s. ("{", "}", "(", ")", "[", "]", "@", "#", ";", ":", "<", ">", "&", '"', "'", "\\")""")
continue
elif any(char in username_input_cleaned for char in banned_chars):
await websocket.send_text(r"""[-] ERROR: Username contains banned character/s. ("{", "}", "(", ")", "[", "]", "@", "#", ";", ":", "<", ">", "&", '"', "'", "\\")""")
continue
elif cleaned_display_name in ["zach(thewebsitecreator)", "zachthewebsitecreator"] or any(item in cleaned_username for item in banned_keywords) or any(item in cleaned_display_name for item in banned_keywords):
await websocket.send_text("[-] ERROR: You cannot impersonate me.")
continue
hashed = await asyncio.to_thread(hash_password, password)
async with file_lock:
assigned_id = str(NEXT_USER_ID)
NEXT_USER_ID += 1
USERS_DB[username_input_cleaned] = {
"user_id": assigned_id,
"password": hashed,
"display_name": display_name,
"friends": [],
"pending_friend_requests": [],
"sent_friend_requests": [],
"session_token": ""
}
async with file_lock:
SALTS_DB[USERS_DB[username_input_cleaned]["user_id"]] = secrets.token_hex(32)
await asyncio.to_thread(save_users_local)
await asyncio.to_thread(save_salts_local)
await websocket.send_text("[+] SUCCESS: Account created! Please log in.")
continue
else:
await websocket.send_text("[-] ERROR: Your password is too short. Please enter a longer password.")
continue
elif action == "token_login":
token_input = auth_data.get("token")
remember_token_input = auth_data.get("remember_token")
matched_user = None
used_remember_token = False
if token_input:
matched_user = ACTIVE_TOKENS.get(token_input)
if not matched_user and remember_token_input:
for u_key, u_info in USERS_DB.items():
u_id = str(u_info.get("user_id"))
u_salt = SALTS_DB.get(u_id)
stored_hash = u_info.get("remember_token_hash")
if u_salt and stored_hash:
check_hash = hashlib.sha256((remember_token_input + u_salt).encode()).hexdigest()
if check_hash == stored_hash:
matched_user = u_key
used_remember_token = True
break
if matched_user:
username = matched_user
if username in ONLINE_USERS:
try:
await ONLINE_USERS[username].close()
except Exception:
pass
ONLINE_USERS[username] = websocket
is_authenticated = True
new_token = str(uuid.uuid4())
if token_input in ACTIVE_TOKENS:
del ACTIVE_TOKENS[token_input]
ACTIVE_TOKENS[new_token] = username
async with file_lock:
USERS_DB[username]["session_token"] = new_token
new_remember_token = None
if used_remember_token:
user_id = str(USERS_DB[username]["user_id"])
user_salt = SALTS_DB.get(user_id)
if user_salt:
new_remember_token = secrets.token_hex(32)
new_combined = new_remember_token + user_salt
new_salted_hash = hashlib.sha256(new_combined.encode()).hexdigest()
async with file_lock:
USERS_DB[username]["remember_token_hash"] = new_salted_hash
await asyncio.to_thread(save_users_local)
token_login_payload = {
"action": "login_success",
"token": new_token,
"user": username,
"display_name": USERS_DB[username]["display_name"],
"pending_friend_requests": USERS_DB[username]["pending_friend_requests"],
"friends": USERS_DB[username]["friends"],
"sent_friend_requests": USERS_DB[username].get("sent_friend_requests", []),
"recent_chats": RECENTS_DB.get(username, [])
}
if new_remember_token:
token_login_payload["remember_token"] = new_remember_token
await websocket.send_text(json.dumps(token_login_payload))
if username in OFFLINE_QUEUES_DB:
for missed_msgs in OFFLINE_QUEUES_DB[username]:
await websocket.send_text(missed_msgs)
del OFFLINE_QUEUES_DB[username]
await asyncio.to_thread(save_offline_queues_local)
print(f"{username} re-authenticated via token.")
else:
await websocket.send_text("[-] FAIL: Invalid or expired session token.")
continue
elif action == "login":
if not username_input_cleaned or not password:
await websocket.send_text("[-] FAIL: Username and password required.")
continue
is_valid = False
if username_input_cleaned in USERS_DB:
is_valid = await asyncio.to_thread(verify_password, password, USERS_DB[username_input_cleaned]["password"])
if is_valid:
username = username_input_cleaned
if username in ONLINE_USERS:
try:
await ONLINE_USERS[username].close()
except Exception:
pass
ONLINE_USERS[username] = websocket
is_authenticated = True
token = str(uuid.uuid4())
ACTIVE_TOKENS[token] = username
user_id = str(USERS_DB[username]["user_id"])
remember_me = auth_data.get("remember_me", False)
remember_token = None
if remember_me:
user_salt = SALTS_DB.get(user_id)
if user_salt:
remember_token = secrets.token_hex(32)
combined_string = remember_token + user_salt
salted_hash = hashlib.sha256(combined_string.encode()).hexdigest()
async with file_lock:
USERS_DB[username]["remember_token_hash"] = salted_hash
await asyncio.to_thread(save_users_local)
async with file_lock:
USERS_DB[username]["session_token"] = token
await asyncio.to_thread(save_users_local)
response_payload = {
"action": "login_success",
"token": token,
"user": username,
"user_id": user_id,
"display_name": USERS_DB[username]["display_name"],
"pending_friend_requests": USERS_DB[username]["pending_friend_requests"],
"friends": USERS_DB[username]["friends"],
"sent_friend_requests": USERS_DB[username].get("sent_friend_requests", []),
"recent_chats": RECENTS_DB.get(username, [])
}
if remember_token:
response_payload["remember_token"] = remember_token
await websocket.send_text(json.dumps(response_payload))
if username in OFFLINE_QUEUES_DB:
for missed_msgs in OFFLINE_QUEUES_DB[username]:
await websocket.send_text(missed_msgs)
del OFFLINE_QUEUES_DB[username]
await asyncio.to_thread(save_offline_queues_local)
print(f"{username} logged in.")
continue
else:
await websocket.send_text("[-] FAIL: Invalid username or password.")
continue
if not is_authenticated:
await websocket.send_text("[-] ERROR: You must log in or sign up first.")
continue
if action == "logout":
async with file_lock:
old_token = USERS_DB[username].get("session_token")
if old_token in ACTIVE_TOKENS:
del ACTIVE_TOKENS[old_token]
USERS_DB[username]["session_token"] = ""
USERS_DB[username]["remember_token_hash"] = ""
ONLINE_USERS.pop(username, None)
await asyncio.to_thread(save_users_local)
await websocket.send_text("[+] SUCCESS: Logged out.")
await websocket.close()
break
# Master list for all the functions
dispatch_table = {
# I won't be adding comments for these because they're just self-explanatory
"search_users": search_users,
"send_friend_request": send_friend_request,
"accept_friend_request": accept_friend_request,
"decline_friend_request": decline_friend_request,
"unfriend": unfriend,
"send_chat_message": send_chat_message,
"request_chat_history": request_chat_history,
"mark_seen": mark_seen,
"ping": ping
}
if action in dispatch_table:
dispatch_table[action]()
else:
await websocket.send_text("[-] ERROR: Invalid action.")
continue
except WebSocketDisconnect:
pass
except Exception as e:
print(f"Unhandled error in chat_handler for {username}: {e}")
finally:
if username and username in ONLINE_USERS and ONLINE_USERS[username] == websocket:
del ONLINE_USERS[username]
print(f"{username} disconnected.")
if __name__ == "__main__":
import uvicorn
print("[+] Starting the messenger backend on port 7860...")
uvicorn.run(app, host="0.0.0.0", port=7860)
the_unreachable_void("Zach", "pass123", "pass123", "Parabola", "No", "NoTokenForYouLol", "pass", "Yes", 1, "Goodbye")