prelude-cli-beta 1446__py3-none-any.whl

prelude_cli_beta/views/detect.py

@@ -0,0 +1,438 @@
+import asyncio
+import click
+import yaml
+
+from datetime import datetime, time, timedelta, timezone
+from dateutil.parser import parse
+from pathlib import Path, PurePath
+
+from prelude_cli_beta.views.shared import Spinner, pretty_print
+from prelude_sdk_beta.controllers.detect_controller import DetectController
+from prelude_sdk_beta.controllers.iam_controller import IAMAccountController
+from prelude_sdk_beta.models.codes import Control, RunCode
+
+
+@click.group()
+@click.pass_context
+def detect(ctx):
+    """Continuous security testing"""
+    ctx.obj = DetectController(account=ctx.obj)
+
+
+@detect.command("create-endpoint")
+@click.option("-h", "--host", help="hostname of this machine", type=str, required=True)
+@click.option(
+    "-s", "--serial_num", help="serial number of this machine", type=str, required=True
+)
+@click.option(
+    "-r",
+    "--reg_string",
+    help="registration string, in the format of <account_id>/<service_user_token>",
+    type=str,
+    required=True,
+)
+@click.option(
+    "-t",
+    "--tags",
+    help="a comma-separated list of tags for this endpoint",
+    type=str,
+    default=None,
+)
+@click.pass_obj
+@pretty_print
+def register_endpoint(controller, host, serial_num, reg_string, tags):
+    """Register a new endpoint"""
+    with Spinner(description="Registering endpoint"):
+        token = controller.register_endpoint(
+            host=host, serial_num=serial_num, reg_string=reg_string, tags=tags
+        )
+    return dict(token=token)
+
+
+@detect.command("update-endpoint")
+@click.argument("endpoint_id")
+@click.option(
+    "-t",
+    "--tags",
+    help="a comma-separated list of tags for this endpoint",
+    type=str,
+    default=None,
+)
+@click.pass_obj
+@pretty_print
+def update_endpoint(controller, endpoint_id, tags):
+    """Update an existing endpoint"""
+    with Spinner(description="Updating endpoint"):
+        return controller.update_endpoint(endpoint_id=endpoint_id, tags=tags)
+
+
+@detect.command("tests")
+@click.option("--techniques", help="comma-separated list of techniques", type=str)
+@click.pass_obj
+@pretty_print
+def list_tests(controller, techniques):
+    """List all security tests"""
+    with Spinner(description="Fetching all security tests"):
+        filters = dict()
+        if techniques:
+            filters["techniques"] = techniques
+        return controller.list_tests(filters=filters)
+
+
+@detect.command("test")
+@click.argument("test_id")
+@click.pass_obj
+@pretty_print
+def get_test(controller, test_id):
+    """List properties for a test"""
+    with Spinner(description="Fetching data for test"):
+        return controller.get_test(test_id=test_id)
+
+
+@detect.command("techniques")
+@click.pass_obj
+@pretty_print
+def list_techniquess(controller):
+    """List techniques"""
+    with Spinner(description="Fetching techniques"):
+        return controller.list_techniques()
+
+
+@detect.command("threats")
+@click.pass_obj
+@pretty_print
+def list_threats(controller):
+    """List all threats"""
+    with Spinner(description="Fetching threats"):
+        return controller.list_threats()
+
+
+@detect.command("threat")
+@click.argument("threat_id")
+@click.pass_obj
+@pretty_print
+def get_threat(controller, threat_id):
+    """List properties for a threat"""
+    with Spinner(description="Fetching data for threat"):
+        return controller.get_threat(threat_id=threat_id)
+
+
+@detect.command("detections")
+@click.pass_obj
+@pretty_print
+def list_detections(controller):
+    """List all Prelude detections"""
+    with Spinner(description="Fetching detections"):
+        return controller.list_detections()
+
+
+@detect.command("detection")
+@click.argument("detection_id")
+@click.option(
+    "-o",
+    "--output_file",
+    help="write the sigma rule to a file in yaml format",
+    type=click.Path(writable=True),
+)
+@click.pass_obj
+@pretty_print
+def get_detection(controller, detection_id, output_file):
+    """List properties for a detection"""
+    with Spinner(description="Fetching data for detection"):
+        data = controller.get_detection(detection_id=detection_id)
+    if output_file:
+        with open(output_file, "w") as f:
+            f.write(yaml.safe_dump(data["rule"]))
+    return data
+
+
+@detect.command("threat-hunts")
+@click.option("--tests", help="comma-separated list of tests", type=str)
+@click.pass_obj
+@pretty_print
+def list_threat_hunts(controller, tests):
+    """List threat hunts"""
+    with Spinner(description="Fetching threat hunts"):
+        filters = dict()
+        if tests:
+            filters["tests"] = tests
+        return controller.list_threat_hunts(filters)
+
+
+@detect.command("threat-hunt")
+@click.argument("threat_hunt_id")
+@click.pass_obj
+@pretty_print
+def get_threat_hunt(controller, threat_hunt_id):
+    """List properties for a threat hunt"""
+    with Spinner(description="Fetching data for threat hunt"):
+        return controller.get_threat_hunt(threat_hunt_id=threat_hunt_id)
+
+
+@detect.command("do-threat-hunt")
+@click.argument("threat_hunt_id")
+@click.pass_obj
+@pretty_print
+def do_threat_hunt(controller, threat_hunt_id):
+    """Run a threat hunt query"""
+    with Spinner(description="Running threat hunt"):
+        return controller.do_threat_hunt(threat_hunt_id=threat_hunt_id)
+
+
+@detect.command("download")
+@click.argument("test")
+@click.pass_obj
+@pretty_print
+def download(controller, test):
+    """Download a test to your local environment"""
+    Path(test).mkdir(parents=True, exist_ok=True)
+    with Spinner(description="Downloading test"):
+        attachments = controller.get_test(test_id=test).get("attachments")
+
+        for attach in attachments:
+            code = controller.download(test_id=test, filename=attach)
+            with open(PurePath(test, attach), "wb") as f:
+                f.write(code)
+
+
+@detect.command("schedule")
+@click.argument("id")
+@click.option(
+    "-t",
+    "--type",
+    help="whether you are scheduling a test or threat",
+    required=True,
+    type=click.Choice(["TEST", "THREAT"], case_sensitive=False),
+)
+@click.option(
+    "--tags",
+    help="only enable for these tags (comma-separated list)",
+    type=str,
+    default="",
+)
+@click.option(
+    "-r",
+    "--run_code",
+    help="provide a run_code",
+    default=RunCode.DAILY.name,
+    show_default=True,
+    type=click.Choice(
+        [r.name for r in RunCode if r != RunCode.INVALID], case_sensitive=False
+    ),
+)
+@click.pass_obj
+@pretty_print
+def schedule(controller, id, type, run_code, tags):
+    """Add test or threat to your queue"""
+    with Spinner(description=f"Scheduling {type.lower()}"):
+        if type == "TEST":
+            return controller.schedule([dict(test_id=id, run_code=run_code, tags=tags)])
+        else:
+            return controller.schedule(
+                [dict(threat_id=id, run_code=run_code, tags=tags)]
+            )
+
+
+@detect.command("unschedule")
+@click.argument("id")
+@click.option(
+    "-t",
+    "--type",
+    help="whether you are unscheduling a test or threat",
+    required=True,
+    type=click.Choice(["TEST", "THREAT"], case_sensitive=False),
+)
+@click.option(
+    "--tags",
+    help="only disable for these tags (comma-separated list)",
+    type=str,
+    default="",
+)
+@click.confirmation_option(prompt="Are you sure?")
+@click.pass_obj
+@pretty_print
+def unschedule(controller, id, type, tags):
+    """Remove test or threat from your queue"""
+    with Spinner(description=f"Unscheduling {type.lower()}"):
+        if type == "TEST":
+            return controller.unschedule([dict(test_id=id, tags=tags)])
+        else:
+            return controller.unschedule([dict(threat_id=id, tags=tags)])
+
+
+@detect.command("delete-endpoint")
+@click.argument("endpoint_id")
+@click.confirmation_option(prompt="Are you sure?")
+@click.pass_obj
+@pretty_print
+def delete_endpoint(controller, endpoint_id):
+    """Delete a probe/endpoint"""
+    with Spinner(description="Deleting endpoint"):
+        return controller.delete_endpoint(ident=endpoint_id)
+
+
+@detect.command("queue")
+@click.pass_obj
+@pretty_print
+def queue(controller):
+    """List all tests in your active queue"""
+    with Spinner(description="Fetching active tests from queue"):
+        iam = IAMAccountController(account=controller.account)
+        return iam.get_account().get("queue")
+
+
+@detect.command("endpoints")
+@click.option(
+    "-d",
+    "--days",
+    help="only show endpoints that have run at least once in the past DAYS days",
+    default=90,
+    type=int,
+)
+@click.pass_obj
+@pretty_print
+def endpoints(controller, days):
+    """List all active endpoints associated to your account"""
+    with Spinner(description="Fetching endpoints"):
+        return controller.list_endpoints(days=days)
+
+
+@detect.command("clone")
+@click.pass_obj
+@pretty_print
+def clone(controller):
+    """Download all tests to your local environment"""
+
+    async def fetch(test):
+        Path(test["id"]).mkdir(parents=True, exist_ok=True)
+
+        for attach in controller.get_test(test_id=test["id"]).get("attachments"):
+            code = controller.download(test_id=test["id"], filename=attach)
+            with open(PurePath(test["id"], attach), "wb") as f:
+                f.write(code)
+
+    async def start_cloning():
+        await asyncio.gather(*[fetch(test) for test in controller.list_tests()])
+
+    with Spinner(description="Downloading all tests"):
+        asyncio.run(start_cloning())
+
+
+@detect.command("activity")
+@click.option(
+    "--view",
+    help="retrieve a specific result view",
+    default="logs",
+    show_default=True,
+    type=click.Choice(
+        [
+            "endpoints",
+            "findings",
+            "logs",
+            "metrics",
+            "protected",
+            "techniques",
+            "tests",
+            "threats",
+        ]
+    ),
+)
+@click.option(
+    "--control", type=click.Choice([c.name for c in Control], case_sensitive=False)
+)
+@click.option("--dos", help="comma-separated list of DOS", type=str)
+@click.option("--endpoints", help="comma-separated list of endpoint IDs", type=str)
+@click.option("--finish", help="end date of activity (end of day)", type=str)
+@click.option("--os", help="comma-separated list of OS", type=str)
+@click.option("--policy", help="comma-separated list of policies", type=str)
+@click.option(
+    "--social",
+    help="whether to fetch account-specific or social stats. Applicable to the following views: protected",
+    is_flag=True,
+)
+@click.option("--start", help="start date of activity (beginning of day)", type=str)
+@click.option("--statuses", help="comma-separated list of statuses", type=str)
+@click.option("--techniques", help="comma-separated list of techniques", type=str)
+@click.option("--tests", help="comma-separated list of test IDs", type=str)
+@click.option("--threats", help="comma-separated list of threat IDs", type=str)
+@click.pass_obj
+@pretty_print
+def describe_activity(
+    controller,
+    control,
+    dos,
+    endpoints,
+    finish,
+    os,
+    policy,
+    social,
+    start,
+    statuses,
+    techniques,
+    tests,
+    threats,
+    view,
+):
+    """View my Detect results"""
+    start = parse(start) if start else datetime.now(timezone.utc) - timedelta(days=29)
+    finish = parse(finish) if finish else datetime.now(timezone.utc)
+    filters = dict(
+        start=datetime.combine(start, time.min),
+        finish=datetime.combine(finish, time.max),
+    )
+    if control:
+        filters["control"] = Control[control.upper()].value
+    if dos:
+        filters["dos"] = dos
+    if endpoints:
+        filters["endpoints"] = endpoints
+    if os:
+        filters["os"] = os
+    if policy:
+        filters["policy"] = policy
+    if social:
+        filters["impersonate"] = "social"
+    if statuses:
+        filters["statuses"] = statuses
+    if techniques:
+        filters["techniques"] = techniques
+    if tests:
+        filters["tests"] = tests
+    if threats:
+        filters["threats"] = threats
+
+    with Spinner(description="Fetching activity"):
+        return controller.describe_activity(view=view, filters=filters)
+
+
+@detect.command("threat-hunt-activity")
+@click.argument("id")
+@click.option(
+    "-t",
+    "--type",
+    help="whether you are getting activity for a threat hunt, test, or threat",
+    required=True,
+    type=click.Choice(["THREAT_HUNT", "TEST", "THREAT"], case_sensitive=False),
+)
+@click.pass_obj
+@pretty_print
+def threat_hunt_activity(controller, id, type):
+    """Get threat hunt activity"""
+    with Spinner(description="Fetching threat hunt activity"):
+        if type == "THREAT_HUNT":
+            return controller.threat_hunt_activity(threat_hunt_id=id)
+        elif type == "TEST":
+            return controller.threat_hunt_activity(test_id=id)
+        else:
+            return controller.threat_hunt_activity(threat_id=id)
+
+
+@detect.command("accept-terms", hidden=True)
+@click.argument("name", type=str, required=True)
+@click.option("-v", "--version", type=str, required=True)
+@click.pass_obj
+@pretty_print
+def accept_terms(controller, name, version):
+    """Accept terms and conditions"""
+    with Spinner(description="Accepting terms and conditions"):
+        return controller.accept_terms(name=name, version=version)

prelude_cli_beta/views/generate.py

@@ -0,0 +1,130 @@
+import json
+import os
+
+import click
+
+from prelude_cli_beta.views.shared import Spinner, pretty_print
+from prelude_sdk_beta.controllers.generate_controller import GenerateController
+from prelude_sdk_beta.models.codes import Control
+
+
+@click.group()
+@click.pass_context
+def generate(ctx):
+    """Generate tests"""
+    ctx.obj = GenerateController(account=ctx.obj)
+
+
+def _process_results(result: dict, output_dir: str, job_id: str) -> dict:
+    if result["status"] == "COMPLETE":
+        for technique in result["output"]:
+            if technique["status"] == "SUCCEEDED":
+                technique_directory = technique["technique"].replace(".", "_")
+                os.makedirs(f"{output_dir}/{technique_directory}")
+                content = technique.get("ai_generated") or technique.get(
+                    "existing_test"
+                )
+                with open(f"{output_dir}/{technique_directory}/test.go", "w") as f:
+                    f.write(content["go_code"])
+                for i, sigma_rule in enumerate(content["sigma_rules"]):
+                    with open(
+                        f"{output_dir}/{technique_directory}/sigma_{i}.yaml", "w"
+                    ) as f:
+                        f.write(sigma_rule)
+                for i, query in enumerate(content.get("threat_hunt_queries", [])):
+                    with open(
+                        f"{output_dir}/{technique_directory}/query_{i}.json", "w"
+                    ) as f:
+                        json.dump(
+                            dict(name=query["name"], query=query["query"]), f, indent=4
+                        )
+                with open(f"{output_dir}/{technique_directory}/config.json", "w") as f:
+                    json.dump(
+                        dict(
+                            technique=technique["technique"],
+                            name=technique["name"],
+                            unit="response",
+                        ),
+                        f,
+                        indent=4,
+                    )
+                if content["readme"]:
+                    with open(
+                        f"{output_dir}/{technique_directory}/README.md", "w"
+                    ) as f:
+                        f.write(content["readme"])
+        return dict(
+            output_dir=output_dir,
+            successfully_generated=[
+                t["technique"] for t in result["output"] if t["status"] == "SUCCEEDED"
+            ],
+            failed=[
+                t["technique"] for t in result["output"] if t["status"] == "FAILED"
+            ],
+            job_id=job_id,
+        )
+    else:
+        raise Exception(
+            "Failed to generate threat intel: %s (ref: %s)", result["reason"], job_id
+        )
+
+
+@generate.command("threat-intel")
+@click.argument(
+    "threat_pdf",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True),
+)
+@click.argument(
+    "output_dir", type=click.Path(dir_okay=True, file_okay=False, writable=True)
+)
+@click.pass_obj
+@pretty_print
+def generate_threat_intel(
+    controller: GenerateController, threat_pdf: str, output_dir: str
+):
+    with Spinner("Uploading") as spinner:
+        job_id = controller.upload_threat_intel(threat_pdf)["job_id"]
+        spinner.update(spinner.task_ids[-1], description="Parsing PDF")
+        while (result := controller.get_threat_intel(job_id)) and result[
+            "status"
+        ] == "RUNNING":
+            if result["step"] == "GENERATE":
+                spinner.update(
+                    spinner.task_ids[-1],
+                    description=f'Generating ({result["completed_tasks"]}/{result["num_tasks"]})',
+                )
+    return _process_results(result, output_dir, job_id)
+
+
+@generate.command("from-advisory")
+@click.argument(
+    "partner", type=click.Choice([Control.CROWDSTRIKE.name], case_sensitive=False)
+)
+@click.option(
+    "-a", "--advisory_id", required=True, type=str, help="Partner advisory ID"
+)
+@click.option(
+    "-o",
+    "--output_dir",
+    required=True,
+    type=click.Path(dir_okay=True, file_okay=False, writable=True),
+)
+@click.pass_obj
+@pretty_print
+def generate_from_partner_advisory(
+    controller: GenerateController, partner: Control, advisory_id: str, output_dir: str
+):
+    with Spinner("Uploading") as spinner:
+        job_id = controller.generate_from_partner_advisory(
+            partner=Control[partner], advisory_id=advisory_id
+        )["job_id"]
+        spinner.update(spinner.task_ids[-1], description="Parsing PDF")
+        while (result := controller.get_threat_intel(job_id)) and result[
+            "status"
+        ] == "RUNNING":
+            if result["step"] == "GENERATE":
+                spinner.update(
+                    spinner.task_ids[-1],
+                    description=f'Generating ({result["completed_tasks"]}/{result["num_tasks"]})',
+                )
+    return _process_results(result, output_dir, job_id)