add: xfarm also like a python module using python -m xfarm

client/exploitfarm/xfarm.py

@@ -0,0 +1,391 @@
+#!/usr/bin/env python3
+
+import typer
+from rich import print
+from rich.markup import escape
+from rich.console import Console
+
+from typer import Abort
+from enum import Enum
+from exploitfarm.utils.reqs import get_url
+from exploitfarm.cmd.config import InitialConfiguration, inital_config_setup, ClientConfig
+from exploitfarm.cmd.login import login_required, try_authenticate
+from exploitfarm.cmd.exploitinit import ExploitConf
+from exploitfarm.utils.config import ExploitConfig, check_exploit_config_exists
+import getpass, re, os, orjson
+from pydantic import PositiveInt
+from typing import Optional
+from uuid import UUID
+from exploitfarm.model import Language
+from exploitfarm.utils.config import EXPLOIT_CONFIG_REGEX
+from exploitfarm.utils import restart_program
+from exploitfarm.cmd.startxploit import start_exploit_tui
+from exploitfarm.utils.reqs import ReqsError
+from requests.exceptions import Timeout as RequestsTimeout
+from exploitfarm import __version__
+import multiprocessing
+from queue import Queue
+
+import traceback
+
+app = typer.Typer(
+    no_args_is_help=True,
+    context_settings={"help_option_names": ["-h", "--help"]}
+)
+console = Console()
+
+DEV_MODE = __version__ == "0.0.0"
+
+class g:
+    interactive = True
+    config: ClientConfig = ClientConfig.read()
+
+def tuple_version(version):
+    return tuple(map(int, version.split(".")))
+
+def initial_setup(login=True):
+    connection = inital_config_setup(g.config, interactive=g.interactive)
+    if g.config.status["version"] != __version__:
+        print("[bold yellow]The server version is different from the client version! This may cause problems![/]")
+        print(f"[bold yellow]Server version: {g.config.status['version']}, Client version: {__version__}[/]")
+        if not typer.confirm("Do you want to continue?", default=False):
+            raise Abort()
+    if DEV_MODE:
+        print("[bold yellow]Development mode detected!")
+    if login and connection:
+        login_required(g.config, interactive=g.interactive)
+    return connection
+
+@app.command(help="Configure the client settings")
+def config(
+    address: str = typer.Option(None, help="The address of the server"),
+    port: int = typer.Option(None, help="The port of the server"),
+    nickname: str = typer.Option(None, help="The nickname of this client"),
+    https: bool = typer.Option(False, help="Use HTTPS for the connection")
+):
+    if g.interactive:
+        init_config = InitialConfiguration(g.config)
+        if init_config.run() == 0:
+            print("[bold green]Configuration saved![/]")
+        else:
+            print("[bold red]Configuration cancelled[/]")
+    else:
+        if address:
+            g.config.server.address = address
+        if port:
+            g.config.server.port = port
+        if nickname:
+            g.config.client_name = nickname
+        if https:
+            g.config.server.https = https
+        elif not g.config.test_server():
+            print(f"[bold red]Connection test failed to {escape(get_url('//', g.config))}[/]")
+            return
+        g.config.write()
+        print("[bold green]Config updated[/]")
+
+@app.command(help="Reset the client settings")
+def reset():
+    print("[bold yellow]Are you sure you want to reset configs?\n[bold red]This operation may break some exploits running on the client.", end="")
+    delete = typer.confirm("")
+    if delete:
+        ClientConfig().write()
+        print("[bold green]Client resetted successful[/]")
+    else:
+        print("[bold]Reset cancelled[/]")
+
+@app.command(help="Start the exploit")
+def start(
+    path: str = typer.Argument(".", help="The path of the exploit"),
+    pool_size: PositiveInt = typer.Option(multiprocessing.cpu_count()*10, "--pool-size", "-p", help="Use fixed thread pool size for the exploit"),
+    submit_pool_timeout: PositiveInt = typer.Option(3, help="The timeout for the submit pool to wait for new attack results and send flags"),
+    server_status_refresh_period: PositiveInt = typer.Option(5, help="The period to refresh the server status"),
+    test: Optional[str] = typer.Option(None, "--test", "-t", help="Test the exploit"),
+    test_timeout: PositiveInt = typer.Option(30, help="The timeout for the test"),
+    max_mem_usage: PositiveInt = typer.Option(95, help="The maximum memory percentage to use of the PC")
+):
+    if max_mem_usage > 100:
+        print("[bold red]Max memory usage can't be greater than 100%[/]")
+        return
+    path = os.path.abspath(path)
+    from exploitfarm.xploit import start_xploit, shutdown, xploit_one
+
+    if not os.path.isdir(path):
+        print(f"[bold red]Path {escape(path)} not found[/]")
+        return
+
+    if not check_exploit_config_exists(path):
+        print(f"[bold red]Exploit configuration not found in {escape(path)}[/]")
+        return
+
+    if test:
+        xploit_one(g.config, test, path, test_timeout)
+        return
+
+    if not initial_setup():
+        print("[bold red]Can't connect to the server! The server is needed to start the exploit! Configure with 'xfarm config'[/]")
+        return
+
+    try:
+        exploit_config = ExploitConfig.read(path)
+        if exploit_config.service not in [UUID(ele["id"]) for ele in g.config.status["services"]]:
+            if not g.interactive:
+                print(f"[bold red]Service {escape(str(exploit_config.service))} not found[/]")
+                return
+            print(f"[bold red]Service {escape(str(exploit_config.service))} not found use 'xfarm init --edit'[/]")
+            decision = typer.confirm("Do you want to continue run 'xfarm init --edit' ?", default=True)
+            if decision:
+                init(edit=True)
+                restart_program()
+            return
+        exploit_config.publish_exploit(g.config)
+    except Exception as e:
+        traceback.print_exc()
+        print(f"[bold red]Error reading exploit configuration from {path}: {e}[/]")
+        return
+
+    shared_infos = {}
+    print_queue = Queue()
+    shared_infos["config"] = g.config.status
+    if not exploit_config.lock_exploit():
+        print("[bold yellow]⚠️ Exploit is already running, do you want to continue? (This process will not be tracked)[/]", end="")
+        cont = typer.confirm("", default=False)
+        if not cont:
+            print("[bold red]Operation cancelled[/]")
+            return
+    exit_event = multiprocessing.Event()
+    restart_event = multiprocessing.Event()
+    start_xploit(g.config, shared_infos, print_queue, pool_size, max_mem_usage, path, submit_pool_timeout, server_status_refresh_period, exit_event, restart_event)
+    if g.interactive:
+        start_exploit_tui(g.config, shared_infos, exploit_config, print_queue, pool_size, exit_event, restart_event)
+        shutdown()
+    else:
+        try:
+            while True:
+                print(print_queue.get())
+        except KeyboardInterrupt:
+            print("[bold yellow]Shutting down the exploit[/]")
+            shutdown()
+    if restart_event.is_set():
+        print("[bold yellow]Restarting the exploit[/]")
+        restart_program()
+
+
+@app.command(help="Login to the server")
+def login(
+    password: str = typer.Option(None, help="The password of the user"),
+    stdin: bool = typer.Option(False, help="Read the password from stdin"),
+):
+    initial_setup(login=False)
+
+    if g.config.status["status"] == "setup":
+        print("[bold red]Please configure the server first[/]")
+        return
+    if g.config.status["loggined"] and not g.config.status["config"]["AUTHENTICATION_REQUIRED"]:
+        print("[bold green]Authentication is not required[/]")
+        return
+    if g.config.status["loggined"]:
+        print("[bold green]Already logged in![/]")
+        return
+
+    if stdin or (not password and not g.interactive):
+        if g.interactive:
+            password = getpass.getpass("Password: ")
+        else:
+            password = input("Password: ")
+        status, error = try_authenticate(password, g.config)
+        if status:
+            print("[bold green]Logged in![/]")
+        else:
+            print(f"[bold red]Error: {escape(error)}[/]")
+        return
+
+    if password:
+        status, error = try_authenticate(password, g.config)
+        if status:
+            print("[bold green]Logged in![/]")
+        else:
+            print(f"[bold red]Error: {escape(error)}[/]")
+        return
+
+    login_required(g.config, interactive=g.interactive)
+
+@app.command(help="Logout from the server")
+def logout():
+    g.config.server.auth_key = None
+    g.config.write()
+    print("[bold red]Logged out[/]")
+
+@app.command(help="Test a submitter")
+def submitter_test(
+    path: str = typer.Argument(help="Submitter python script"),
+    kwargs: str = typer.Option("{}", help="Submitter key-words args (json)"),
+    output: str = typer.Argument(help="Text containing flags according to server REGEX")
+):
+    initial_setup()
+    try:
+        kwargs = orjson.loads(kwargs)
+    except Exception as e:
+        print(f"[bold red]Invalid kwargs json: {e}")
+        return
+
+    try:
+        with open(path, "rt") as f:
+            submitter_code = f.read()
+    except Exception as e:
+        print(f"[bold red]File {escape(path)} not found: {e}")
+        return
+
+    if not output:
+        print("[bold red]Output can't be empty")
+
+    flags = [output]
+    if g.config.status["config"]["FLAG_REGEX"]:
+        flags = re.findall(g.config.status["config"]["FLAG_REGEX"], output)
+
+    if len(flags) == 0:
+        print(f"[bold red]No flags extracted from output! REGEX: {escape(g.config.status['config']['FLAG_REGEX'])}")
+        return
+    submitter_id = None
+    try:
+        submitter_id:int = g.config.reqs.new_submitter({
+            "name": "TEST_SUBMITTER (Will be deleted soon)",
+            "kargs": kwargs,
+            "code": submitter_code
+        })["id"]
+        print("[bold yellow]----- TEST RESULTS -----")
+        print("[bold yellow]Flags to submit:[/]", flags)
+        print("[bold yellow]Output:[/]")
+        print(g.config.reqs.test_submitter(submitter_id, flags))
+        print("[bold yellow]----- TEST RESULTS -----")
+    finally:
+        if submitter_id:
+            g.config.reqs.delete_submitter(submitter_id)
+
+class StatusWhat(Enum):
+    status = "status"
+    submiters = "submitters"
+    services = "services"
+    exploits = "exploits"
+    flags = "flags"
+    teams = "teams"
+    clients = "clients"
+
+@app.command(help="Get status of the server")
+def status(
+    what:StatusWhat = typer.Argument(StatusWhat.status.value, help="Server informations type")
+):
+    initial_setup()
+    match what:
+        case StatusWhat.status:
+            print(g.config.status)
+        case StatusWhat.submiters:
+            print(g.config.reqs.submitters())
+        case StatusWhat.services:
+            print(g.config.reqs.services())
+        case StatusWhat.exploits:
+            print(g.config.reqs.exploits())
+        case StatusWhat.flags:
+            print(g.config.reqs.flags())
+        case StatusWhat.teams:
+            print(g.config.reqs.teams())
+        case StatusWhat.clients:
+            print(g.config.reqs.clients())
+
+@app.command(help="Initiate a new exploit configuration")
+def init(
+    edit: bool = typer.Option(False, "--edit", "-e", help="Edit the exploit configuration"),
+    name: Optional[str] = typer.Option(None, help="The name of the exploit"),
+    service: Optional[UUID] = typer.Option(None, help="The service of the exploit"),
+    language: Optional[Language] = typer.Option(None, help="The language of the exploit"),
+):
+    initial_setup()
+    if g.interactive:
+        if edit:
+            if check_exploit_config_exists("."):
+                expl_conf = ExploitConfig.read(".")
+                name = expl_conf.name
+                service = expl_conf.service
+                language = expl_conf.language
+            else:
+                print("[bold red]Exploit configuration not found![/]")
+                return
+            init_config = ExploitConf(g.config, edit, name, service, language)
+        else:
+            init_config = ExploitConf(g.config, edit)
+        final_status = init_config.run()
+        if final_status == 0:
+            print(f"[bold green]Exploit configuration {'created' if not edit else 'edited'}![/]")
+        elif final_status == 99:
+            print("[bold yellow]Exploit folder created, but not registered on the server![/]")
+        else:
+            print("[bold red]Exploit configuration cancelled[/]")
+    else:
+        exists = check_exploit_config_exists(name if not edit else ".")
+
+        if edit ^ exists:
+            print(f"[bold red]Exploit '{escape(name)}' already exists!")
+            return
+
+        if (not name or edit) or not re.match(EXPLOIT_CONFIG_REGEX, name):
+            print(f"[bold red]Please provide a valid name for the exploit (regex: {escape(EXPLOIT_CONFIG_REGEX)})[/]")
+            return
+
+        try:
+            if not service in [UUID(ele["id"]) for ele in g.config.status["services"]]:
+                service = None
+                if not edit:
+                    print("[bold red]Service not found, add a new on the server[/]")
+                    return
+        except Exception:
+            print("[bold red]Service id not found[/]")
+            return
+
+        if (not language or edit):
+            print("[bold red]Language not found[/]")
+            return
+
+        if edit:
+            expl_conf = ExploitConfig.read(name)
+            if name: expl_conf.name = name
+            if language: expl_conf.language = language
+            if service: expl_conf.service = service
+        else:
+            expl_conf = ExploitConfig.new(name, language, service)
+        expl_conf.write(name)
+        expl_conf.publish_exploit(g.config)
+        if edit:
+            print("[bold green]Exploit configuration updated![/]")
+        else:
+            print("[bold green]Exploit configuration created![/]")
+
+def version_callback(verison: bool):
+    if verison:
+        print(__version__, "Development Mode" if DEV_MODE else "Release")
+        raise typer.Exit()
+
+def help_callback(help: bool):
+    if help:
+        raise typer.Exit()
+
+@app.callback()
+def main(
+    no_interactive: bool = typer.Option(False, "--no-interactive", "-I", help="Interactive configuration mode", envvar="XFARM_INTERACTIVE"),
+    verison: bool = typer.Option(False, "--version", "-v", help="Show the version of the client", callback=version_callback),
+):
+    g.interactive = not no_interactive
+
+def run():
+    try:
+        app()
+    except KeyboardInterrupt:
+        print("[bold yellow]Operation cancelled[/]")
+    except Abort:
+        print("[bold yellow]Operation cancelled[/]")
+    except ReqsError as e:
+        print("[bold red]The server returned an error: {e}[/]")
+    except RequestsTimeout as e:
+        print(f"[bold red]The server has timed out: {e}[/]")
+
+if __name__ == "__main__":
+    run()