import asyncio
import sys
import os
import ssl
import shutil
import tty
import termios
import base64
import argparse
import uuid
import re
import shlex
from websockets.asyncio.client import connect
from websockets.exceptions import ConnectionClosed

DEFAULT_URI = "wss://figure_it_out.cohort.htb/terminal/ws"

async def read_line_locally(reader_stream):
    """Reads a line of text locally while the terminal is still in raw mode, handling backspaces."""
    line_buffer = bytearray()
    while True:
        char = await reader_stream.read(1)
        if not char:
            break

        if char == b'\x03':
            raise KeyboardInterrupt

        # Handle Backspace / Delete
        if char in (b'\x7f', b'\x08'):
            if len(line_buffer) > 0:
                line_buffer.pop()
                # Visually erase character from local terminal screen
                sys.stdout.write('\b \b')
                sys.stdout.flush()
            continue
            
        # Handle Enter key (Carriage Return or Newline)
        if char in (b'\r', b'\n'):
            sys.stdout.write('\r\n')
            sys.stdout.flush()
            break
            
        # Echo the typed character locally so the user can see what they type
        sys.stdout.write(char.decode('utf-8', errors='replace'))
        sys.stdout.flush()
        line_buffer.extend(char)
        
    return line_buffer.decode('utf-8', errors='replace').strip()

async def perform_upload(ws, local_path_raw, remote_path_raw):
    """Reads a local file, base64-encodes it, and pushes a decode command to the remote shell."""
    local_path = os.path.abspath(os.path.expanduser(local_path_raw))

    if not local_path_raw or not os.path.isfile(local_path):
        sys.stdout.write(f"[-] Error: Local file '{local_path}' does not exist.\r\n")
        return

    remote_path = remote_path_raw or os.path.basename(local_path)

    sys.stdout.write(f"[*] Reading and encoding: {local_path}\r\n")
    sys.stdout.flush()

    CHUNK_SIZE = 64 * 1024  # base64 chars per remote shell command, well under typical ARG_MAX

    try:
        with open(local_path, "rb") as f:
            file_bytes = f.read()

        b64_data = base64.b64encode(file_bytes).decode('utf-8')
        quoted_remote_path = shlex.quote(remote_path)
        chunks = [b64_data[i:i + CHUNK_SIZE] for i in range(0, len(b64_data), CHUNK_SIZE)] or [""]

        if len(chunks) > 1:
            sys.stdout.write(f"[*] Pushing encoded payload to target in {len(chunks)} chunks... \r\n")
        else:
            sys.stdout.write(f"[*] Pushing encoded payload to target... \r\n")
        sys.stdout.flush()

        for i, chunk in enumerate(chunks):
            redirect = ">" if i == 0 else ">>"
            await ws.send(f"echo '{chunk}' | base64 -d {redirect} {quoted_remote_path}\n")

        sys.stdout.write("[+] Upload command dispatched seamlessly!\r\n")
    except Exception as e:
        sys.stdout.write(f"[-] Upload failed: {e}\r\n")
    sys.stdout.flush()

async def perform_download(ws, state, remote_path_raw, local_path_raw):
    """Requests a base64 dump of a remote file, captures it off the reader stream, and writes it locally."""
    if not remote_path_raw or not local_path_raw:
        sys.stdout.write("[-] Usage: !download <remote_path> <local_path>\r\n")
        sys.stdout.flush()
        return

    local_path = os.path.abspath(os.path.expanduser(local_path_raw))
    start_tag = f"__DL_START_{uuid.uuid4().hex}__"
    end_tag = f"__DL_END_{uuid.uuid4().hex}__"
    # The remote pty echoes back the literal command text before executing it, so a
    # plain literal marker would false-match on that echo instead of the real output.
    # Suffixing with $$ (shell-expanded PID) means the echo contains the literal,
    # unexpanded "$$" text while only the real output contains digits - so a
    # digit-suffix regex only ever matches the genuine output.
    start_pattern = re.compile(re.escape(start_tag) + r"-\d+")
    end_pattern = re.compile(re.escape(end_tag) + r"-\d+")

    loop = asyncio.get_running_loop()
    state.capturing = True
    state.buffer = ""
    state.end_pattern = end_pattern
    state.future = loop.create_future()

    sys.stdout.write(f"[*] Requesting remote file: {remote_path_raw}\r\n")
    sys.stdout.flush()

    try:
        quoted_remote_path = shlex.quote(remote_path_raw)
        await ws.send(f"echo {start_tag}-$$; base64 {quoted_remote_path} 2>/dev/null; echo {end_tag}-$$\n")
        captured = await asyncio.wait_for(state.future, timeout=30)

        start_match = start_pattern.search(captured)
        end_match = end_pattern.search(captured)
        if not start_match or not end_match:
            sys.stdout.write("[-] Download failed: markers not found in remote response.\r\n")
            return

        payload = captured[start_match.end():end_match.start()].strip()
        if not payload:
            sys.stdout.write(f"[-] Download failed: remote file '{remote_path_raw}' is empty or does not exist.\r\n")
            return

        file_bytes = base64.b64decode(payload, validate=False)
        with open(local_path, "wb") as f:
            f.write(file_bytes)
        sys.stdout.write(f"[+] Downloaded {remote_path_raw} -> {local_path} ({len(file_bytes)} bytes)\r\n")
    except asyncio.TimeoutError:
        sys.stdout.write("[-] Download failed: timed out waiting for remote response.\r\n")
    except Exception as e:
        sys.stdout.write(f"[-] Download failed: {e}\r\n")
    finally:
        state.capturing = False
        state.future = None
        sys.stdout.flush()

async def handle_local_upload_flow(ws, reader_stream):
    """Executes the file upload menu completely locally, hiding input from the remote machine."""
    sys.stdout.write("\r\n" + "="*40 + "\r\n")
    sys.stdout.write("[*] LOCAL FILE UPLOAD INTERCONNECT ACTIVATED\r\n")
    sys.stdout.write("="*40 + "\r\n")
    sys.stdout.write("Enter LOCAL file path: ")
    sys.stdout.flush()

    local_path_raw = await read_line_locally(reader_stream)

    sys.stdout.write("Enter REMOTE destination filename/path [Default: same name]: ")
    sys.stdout.flush()
    remote_path_raw = await read_line_locally(reader_stream)

    await perform_upload(ws, local_path_raw, remote_path_raw)

    sys.stdout.write("="*40 + "\r\n")
    sys.stdout.flush()

class TransferState:
    """Shared state letting writer() ask reader() to capture output instead of printing it (used by !download)."""
    def __init__(self):
        self.capturing = False
        self.buffer = ""
        self.end_pattern = None
        self.future = None

async def reader(ws, state):
    """Blazing fast inbound stream parsing."""
    try:
        async for msg in ws:
            output = msg if isinstance(msg, str) else msg.decode('utf-8', errors='replace')

            if state.capturing:
                state.buffer += output
                # end_pattern only matches the shell-expanded (real) marker, never
                # the literal command text echoed back by the remote pty before it runs.
                if state.end_pattern and state.end_pattern.search(state.buffer):
                    state.capturing = False
                    if state.future and not state.future.done():
                        state.future.set_result(state.buffer)
                continue

            sys.stdout.write(output)
            sys.stdout.flush()
    except ConnectionClosed:
        pass
    except Exception as e:
        print(f"\r\n[Reader error: {e}]")

# ".upload" fires immediately once typed (opens the interactive wizard, no args on the line).
IMMEDIATE_TRIGGERS = (".upload",)
# "!upload"/"!download" are full-line commands: args are parsed from the whole line at Enter.
LINE_TRIGGERS = ("!upload", "!download")

async def writer(ws, old_settings, state):
    """Forwards keystrokes to the remote instantly (so arrow-key history / tab-completion work),
    except while a line might still turn into a local trigger command, which is held back and
    edited locally until it's ruled out or completed."""
    loop = asyncio.get_running_loop()
    reader_stream = asyncio.StreamReader()
    protocol = asyncio.StreamReaderProtocol(reader_stream)

    await loop.connect_read_pipe(lambda: protocol, sys.stdin)

    input_buffer = ""   # full text of the current line, used for /exit /bye matching
    held_buffer = ""    # chars withheld from the remote while they might still complete a trigger
    checking_trigger = True

    def reset_line():
        nonlocal input_buffer, checking_trigger
        input_buffer = ""
        checking_trigger = True

    try:
        while True:
            # Read character by character (or small paste chunks) instantly
            data = await reader_stream.read(1)
            if not data:
                break

            if b'\x03' in data:
                raise KeyboardInterrupt

            text_data = data.decode('utf-8', errors='replace')

            # Backspace/Delete while composing a possible trigger:  these chars were never sent
            # to the remote, so there's nothing for it to erase - handle editing locally.
            if checking_trigger and data in (b'\x7f', b'\x08'):
                if held_buffer:
                    held_buffer = held_buffer[:-1]
                    input_buffer = input_buffer[:-1]
                    sys.stdout.write('\b \b')
                    sys.stdout.flush()
                continue

            input_buffer += text_data

            if checking_trigger:
                held_buffer += text_data
                has_newline = '\r' in text_data or '\n' in text_data

                # INSTANT INTERCEPT:  Nothing was forwarded to the remote, so nothing to erase/leak.
                if held_buffer == ".upload":
                    sys.stdout.write(text_data)
                    sys.stdout.flush()
                    await handle_local_upload_flow(ws, reader_stream)
                    held_buffer = ""
                    reset_line()
                    continue

                line_trigger = next(
                    (t for t in LINE_TRIGGERS if held_buffer == t or held_buffer.startswith(t + " ")),
                    None,
                )

                if has_newline:
                    if line_trigger:
                        # Move to a fresh line locally - the remote never receives/echoes this line.
                        sys.stdout.write('\r\n')
                        sys.stdout.flush()
                        line = held_buffer.replace('\r', '\n').strip()
                        args = line[len(line_trigger):].strip().split(maxsplit=1)
                        if len(args) != 2:
                            sys.stdout.write(f"[-] Usage: {line_trigger} <arg1> <arg2>\r\n")
                            sys.stdout.flush()
                        elif line_trigger == "!upload":
                            await perform_upload(ws, args[0], args[1])
                        else:
                            await perform_download(ws, state, args[0], args[1])
                        held_buffer = ""
                        reset_line()
                        continue
                    # Not a trigger line after all - erase our local echo of it (all but this
                    # newline char, which was never echoed) and release the whole line to the
                    # remote as one chunk so its own pty echo redisplays it correctly.
                    checking_trigger = False
                    erase_count = len(held_buffer) - 1
                    if erase_count > 0:
                        sys.stdout.write('\b \b' * erase_count)
                        sys.stdout.flush()
                    flushed = held_buffer
                    held_buffer = ""
                    await ws.send(flushed)
                else:
                    could_still_match = (
                        any(t.startswith(held_buffer) for t in IMMEDIATE_TRIGGERS)
                        or any(t.startswith(held_buffer) or held_buffer.startswith(t) for t in LINE_TRIGGERS)
                    )
                    if could_still_match:
                        # Could still become a trigger (word itself, or its args) - hold it back
                        # from the remote, but echo it locally so the user can see what they typed.
                        sys.stdout.write(text_data)
                        sys.stdout.flush()
                        continue
                    # Ruled out as a trigger: erase our local echo of the earlier held chars (this
                    # one was never echoed) and release the withheld chars to the remote as one chunk.
                    checking_trigger = False
                    erase_count = len(held_buffer) - 1
                    if erase_count > 0:
                        sys.stdout.write('\b \b' * erase_count)
                        sys.stdout.flush()
                    flushed = held_buffer
                    held_buffer = ""
                    await ws.send(flushed)
                    continue

            else:
                # Steady state for the rest of this line: forward each char immediately, giving
                # the remote every keystroke live (arrow-key history, tab-completion, Ctrl+R, etc).
                await ws.send(text_data)

            # If they hit a newline, reset our command intercept tracking buffer
            if '\r' in text_data or '\n' in text_data:
                clean_command = input_buffer.replace('\r', '\n').strip()

                # Check for the custom local-exit shortcuts (distinct from the remote's own
                # exit/logout, which are no longer blocked and now reach the shell normally).
                if clean_command in ['/exit', '/bye']:
                    print("\r\n[*] Custom exit command recognized. Initiating teardown...")
                    break

                reset_line()

    except ConnectionClosed:
        pass
    except (asyncio.CancelledError, KeyboardInterrupt):
        pass
    except Exception as e:
        print(f"\r\n[Writer error: {e}]")

async def main(target_url):
    ssl_context = ssl.create_default_context()
    ssl_context.check_hostname = False
    ssl_context.verify_mode = ssl.CERT_NONE

    print(f"[*] Connecting to {target_url}...")
    try:
        async with connect(target_url, ssl=ssl_context) as ws:
            print("[+] Connected! Low-latency raw mode active.\r")
            print("="*shutil.get_terminal_size().columns + "\r")
            
            old_settings = termios.tcgetattr(sys.stdin.fileno())
            tty.setraw(sys.stdin.fileno())
            
            state = TransferState()
            try:
                read_task = asyncio.create_task(reader(ws, state))
                write_task = asyncio.create_task(writer(ws, old_settings, state))
                
                done, pending = await asyncio.wait(
                    [read_task, write_task],
                    return_when=asyncio.FIRST_COMPLETED
                )
                
                for task in pending:
                    task.cancel()
            finally:
                termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN, old_settings)
                
    except Exception as e:
        print(f"[-] Connection failed: {e}", file=sys.stderr)
    finally:
        print("\r\n[*] Session terminated.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Low-latency raw websocket terminal client.")
    parser.add_argument(
        "-t",
        "--target_url",
        type=str,
        required=True,
        help=f"Target WebSocket URI (required). For example {DEFAULT_URI}"
    )
    
    args = parser.parse_args()

    try:
        asyncio.run(main(args.target_url))
    except KeyboardInterrupt:
        print("\r\n[*] Exiting via user interrupt.")
        sys.exit(0)