import asyncio
import json
import os
import shlex
import subprocess
import websockets


# =========================
# CONFIG
# =========================

SERVER_URL = "wss://env.neome.com"

WORK_DIR = "."
TERMINAL_PASSWORD = "CHANGE_ME"

WS_MAX_SIZE = 50 * 1024 * 1024
COMMAND_TIMEOUT = 30

# Only commands starting with these names are allowed.
# Add/remove commands here.
ALLOWED_COMMANDS = [
	"ls",
	"dir",
	"pwd",
	"cd",
	"nohup",
	"whoami",
	"date",
	"echo",
	"cat",
	"type",
	"more",
	"python",
	"python3",
	"pip",
	"pip3",
	"git",
	"node",
	"npm",
	"bash",
	"sh",
	"cmd"
]


# =========================
# TERMINAL
# =========================

def run_cmd(command):
	command = str(command or "").strip()

	if not command:
		return "Missing command"

	try:
		parts = shlex.split(command)
	except Exception as e:
		return "Invalid command: " + str(e)

	if not parts:
		return "Missing command"

	is_background = parts[0].lower() == "nohup"

	if is_background:
		parts = parts[1:]

		if not parts:
			return "Missing command after nohup"

	base_cmd = parts[0].lower()

	if base_cmd not in ALLOWED_COMMANDS:
		return (
			"Command not allowed: " +
			base_cmd +
			"\nAllowed: " +
			", ".join(ALLOWED_COMMANDS)
		)

	try:
		if is_background:
			p = subprocess.Popen(
				parts,
				cwd=WORK_DIR,
				text=True,
				stdout=subprocess.DEVNULL,
				stderr=subprocess.DEVNULL,
				start_new_session=True
			)

			return "Started background process PID " + str(p.pid)

		result = subprocess.run(
			parts,
			cwd=WORK_DIR,
			text=True,
			capture_output=True,
			timeout=COMMAND_TIMEOUT
		)

		out = ""

		if result.stdout:
			out += result.stdout

		if result.stderr:
			if out:
				out += "\n"
			out += result.stderr

		if result.returncode != 0:
			out = "Exit code: " + str(result.returncode) + "\n" + out

		if not out.strip():
			out = "OK"

		return out.strip()

	except subprocess.TimeoutExpired as e:
		out = ""

		if e.stdout:
			out += str(e.stdout)

		if e.stderr:
			if out:
				out += "\n"
			out += str(e.stderr)

		out = out.strip()

		if out:
			return out + "\n\nError: command timeout after " + str(COMMAND_TIMEOUT) + " seconds"

		return "Error: command timeout after " + str(COMMAND_TIMEOUT) + " seconds"

	except Exception as e:
		return "Error: " + str(e)


# =========================
# COMMANDS
# =========================

def handle_command(cmd):
	try:
		if not isinstance(cmd, list) or not cmd:
			return "Invalid command"

		cmd_type = str(cmd[0] or "").lower()

		if cmd_type == "cmd":
			if len(cmd) < 2:
				return "Missing command"

			return run_cmd(str(cmd[1]))

		return "Unknown command: " + cmd_type

	except Exception as e:
		return "Error: " + str(e)


# =========================
# SOCKET
# =========================

async def send_json(ws, data):
	await ws.send(json.dumps(data))


async def main():
	while True:
		try:
			async with websockets.connect(
				SERVER_URL,
				max_size=WS_MAX_SIZE
			) as ws:

				await send_json(ws, {
					"role": "client"
				})

				print("[+] connected", flush=True)

				async for raw in ws:
					try:
						data = json.loads(raw)
					except:
						continue

					if data.get("type") == "hello":
						print("CLIENT_ID:", data.get("id"), flush=True)
						continue

					if data.get("type") == "cmd":
						result = handle_command(
							data.get("cmd")
						)

						await send_json(ws, {
							"type": "result",
							"output": result
						})

		except Exception as e:
			print("[!] disconnected:", e, flush=True)
			await asyncio.sleep(2)


if __name__ == "__main__":
	asyncio.run(main())