ENTIRELY REWORK IT
This commit is contained in:
213
cvmsentry.py
213
cvmsentry.py
@@ -1,71 +1,18 @@
|
||||
from typing import List, Optional
|
||||
from enum import IntEnum
|
||||
from cvmlib import guac_decode, guac_encode
|
||||
import config
|
||||
import os
|
||||
import websocket
|
||||
import websockets, asyncio
|
||||
from websockets import Subprotocol, Origin
|
||||
import logging
|
||||
import sys
|
||||
import rel
|
||||
import threading
|
||||
import multiprocessing
|
||||
import argparse
|
||||
from datetime import datetime, timezone
|
||||
import json
|
||||
|
||||
LOG_LEVEL = getattr(config, 'log_level', 'INFO')
|
||||
user_state = 0 # 0 = disconnected, 1 = connected, 2 = logged in
|
||||
LOG_LEVEL = getattr(config, "log_level", "INFO")
|
||||
|
||||
def full_login(ws: websocket.WebSocketApp):
|
||||
# Connect
|
||||
vm_name = os.path.basename(ws.url)
|
||||
ws.send(guac_encode("connect", vm_name))
|
||||
ws.send(guac_encode("login", config.credentials["session_auth"]))
|
||||
|
||||
def on_message(ws, message: str):
|
||||
global user_state
|
||||
if guac_decode(message) is not None:
|
||||
decoded: Optional[List[str]] = guac_decode(message)
|
||||
match decoded:
|
||||
case ["nop"]:
|
||||
ws.send(guac_encode("nop"))
|
||||
case ['connect', '1', '1', '1', '0']:
|
||||
print("Connected")
|
||||
user_state = 1
|
||||
case ["auth", "https://auth.collabvm.org"]:
|
||||
full_login(ws)
|
||||
case ["login", "1"]:
|
||||
print("Logged in")
|
||||
user_state = 2
|
||||
case _: # For the XTREME!!!
|
||||
if decoded is not None:
|
||||
if decoded[0] == "sync" or decoded[0] == "png":
|
||||
return
|
||||
elif decoded[0] == "adduser":
|
||||
number_of_users = int(decoded[1])
|
||||
|
||||
if user_state == 0 and number_of_users == 1:
|
||||
return # Ignore ourselves...
|
||||
|
||||
# print(f"Number of users: {number_of_users}")
|
||||
# users = decoded[2:]
|
||||
# print(f"Users: {users}")
|
||||
elif decoded[0] == "chat":
|
||||
user = "System" if len(decoded[1]) == 0 else decoded[1]
|
||||
message = decoded[2]
|
||||
log.info(f"[{user}]: {message}")
|
||||
|
||||
# print(f"Received: {decoded}")
|
||||
|
||||
def on_error(ws, error):
|
||||
|
||||
print(error)
|
||||
|
||||
def on_close(ws, close_status_code, close_msg):
|
||||
print("### closed ###")
|
||||
|
||||
def on_open(ws):
|
||||
ws.send(guac_encode("rename",""))
|
||||
# websocket.enableTrace(True)
|
||||
|
||||
# Prepare logs
|
||||
# Prepare logs
|
||||
if not os.path.exists("logs"):
|
||||
os.makedirs("logs")
|
||||
log_format = logging.Formatter(
|
||||
@@ -79,23 +26,133 @@ log.addHandler(stdout_handler)
|
||||
|
||||
log.info(f"CVM-Sentry started")
|
||||
|
||||
# Parse the command-line argument
|
||||
parser = argparse.ArgumentParser(description="CVM-Sentry")
|
||||
parser.add_argument("vm_name", type=str, help="Name of the VM to connect to (e.g., vm7)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Ensure the provided VM name exists in the config
|
||||
if args.vm_name not in config.vms:
|
||||
log.error(f"VM '{args.vm_name}' not found in configuration.")
|
||||
sys.exit(1)
|
||||
class CollabVMState(IntEnum):
|
||||
WS_CONNECTED = 0
|
||||
VM_CONNECTED = 1
|
||||
LOGGED_IN = 2
|
||||
|
||||
# Use the provided VM name to connect
|
||||
ws = websocket.WebSocketApp(config.vms[args.vm_name],
|
||||
on_open=on_open,
|
||||
on_message=on_message,
|
||||
on_error=on_error,
|
||||
on_close=on_close,
|
||||
header={"User-Agent": "cvmsentry"}, subprotocols=["guacamole"])
|
||||
ws.run_forever(dispatcher=rel, reconnect=5) # Set dispatcher to automatic reconnection, 5 second reconnect delay if connection closed unexpectedly
|
||||
rel.signal(2, rel.abort) # Keyboard Interrupt
|
||||
rel.dispatch()
|
||||
class CollabVMRank(IntEnum):
|
||||
UNREGISTERED = 0
|
||||
REGISTERED = 1
|
||||
ADMIN = 2
|
||||
MOD = 3
|
||||
|
||||
users = {}
|
||||
vm_botuser = {}
|
||||
|
||||
def get_origin_from_ws_url(ws_url: str) -> str:
|
||||
domain = (
|
||||
ws_url.removeprefix("ws:")
|
||||
.removeprefix("wss:")
|
||||
.removeprefix("/")
|
||||
.removeprefix("/")
|
||||
.split("/", 1)[0]
|
||||
)
|
||||
is_wss = ws_url.startswith("wss:")
|
||||
return f"http{'s' if is_wss else ''}://{domain}/"
|
||||
|
||||
|
||||
async def connect(vm_name: str):
|
||||
if vm_name not in config.vms:
|
||||
log.error(f"VM '{vm_name}' not found in configuration.")
|
||||
return
|
||||
uri = config.vms[vm_name]
|
||||
STATE = None
|
||||
log_file_path = os.path.join(getattr(config, "log_directory", "logs"), f"{vm_name}.json")
|
||||
if not os.path.exists(log_file_path):
|
||||
with open(log_file_path, "w") as log_file:
|
||||
log_file.write("{}")
|
||||
async with websockets.connect(
|
||||
uri=uri,
|
||||
subprotocols=[Subprotocol("guacamole")],
|
||||
origin=Origin(get_origin_from_ws_url(uri)),
|
||||
) as websocket:
|
||||
log.info(f"Connected to VM '{vm_name}' at {uri}")
|
||||
STATE = CollabVMState.WS_CONNECTED
|
||||
await websocket.send(guac_encode("rename", ""))
|
||||
await websocket.send(guac_encode("connect", vm_name))
|
||||
if vm_name not in users:
|
||||
users[vm_name] = {}
|
||||
# response = await websocket.recv()
|
||||
async for message in websocket:
|
||||
decoded: Optional[List[str]] = guac_decode(str(message))
|
||||
match decoded:
|
||||
case ["nop"]:
|
||||
await websocket.send(guac_encode("nop"))
|
||||
log.debug((f"({CollabVMState(STATE).name}) Received: {decoded}"))
|
||||
case ["auth", "https://auth.collabvm.org"]:
|
||||
await asyncio.sleep(1)
|
||||
await websocket.send(
|
||||
guac_encode("login", config.credentials["session_auth"])
|
||||
)
|
||||
case ["connect", "1", "1", "1", "0"]:
|
||||
STATE = CollabVMState.VM_CONNECTED
|
||||
log.debug((f"({CollabVMState(STATE).name} - {vm_name}) Connected"))
|
||||
case ["login", "1"]:
|
||||
STATE = CollabVMState.LOGGED_IN
|
||||
case _:
|
||||
if decoded is not None:
|
||||
if decoded[0] in ("sync", "png", "flag", "turn", "size"):
|
||||
continue
|
||||
elif decoded[0] == "chat":
|
||||
user = "System" if len(decoded[1]) == 0 else decoded[1]
|
||||
if user == "System" or STATE < CollabVMState.LOGGED_IN:
|
||||
continue
|
||||
message = decoded[2]
|
||||
log.info(f"[{vm_name} - {user}]: {message}")
|
||||
utc_now = datetime.now(timezone.utc)
|
||||
utc_day = utc_now.strftime("%Y-%m-%d")
|
||||
timestamp = utc_now.isoformat()
|
||||
|
||||
with open(log_file_path, "r+") as log_file:
|
||||
try:
|
||||
log_data = json.load(log_file)
|
||||
except json.JSONDecodeError:
|
||||
log_data = {}
|
||||
|
||||
if utc_day not in log_data:
|
||||
log_data[utc_day] = []
|
||||
|
||||
log_data[utc_day].append({
|
||||
"timestamp": timestamp,
|
||||
"username": user,
|
||||
"message": message
|
||||
})
|
||||
|
||||
log_file.seek(0)
|
||||
json.dump(log_data, log_file, indent=4)
|
||||
log_file.truncate()
|
||||
continue
|
||||
elif decoded[0] == "adduser":
|
||||
if STATE == CollabVMState.LOGGED_IN:
|
||||
username = decoded[2]
|
||||
rank = CollabVMRank(int(decoded[3]))
|
||||
users[vm_name][username] = rank
|
||||
elif STATE < CollabVMState.LOGGED_IN:
|
||||
initial_user_payload = decoded[2:]
|
||||
for i in range(0, len(initial_user_payload), 2):
|
||||
username = initial_user_payload[i]
|
||||
rank = CollabVMRank(int(initial_user_payload[i + 1]))
|
||||
users[vm_name][username] = rank
|
||||
elif decoded[0] == "remuser":
|
||||
if STATE == CollabVMState.LOGGED_IN:
|
||||
username = decoded[2]
|
||||
if username in users[vm_name]:
|
||||
del users[vm_name][username]
|
||||
log.debug(
|
||||
(
|
||||
f"({CollabVMState(STATE).name} - {vm_name}) Received: {decoded}"
|
||||
)
|
||||
)
|
||||
|
||||
for vm in config.vms.keys():
|
||||
|
||||
|
||||
def start_vm_thread(vm_name: str):
|
||||
asyncio.run(connect(vm_name))
|
||||
async def main():
|
||||
tasks = [connect(vm) for vm in config.vms.keys()]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user