Added remaining playbook endpoint integration

pyintelowl/cli/__init__.py

@@ -3,13 +3,15 @@
 from .config import config
 from .investigations import investigations
 from .jobs import jobs
+from .playbooks import playbooks
 from .tags import tags
 
 groups = [
     analyse,
     config,
     jobs,
     tags,
+    playbooks,
     investigations,
 ]
 

pyintelowl/cli/playbooks.py

@@ -0,0 +1,126 @@
+import json
+
+import click
+from rich import box
+from rich import print as rprint
+from rich.console import Console, Group
+from rich.panel import Panel
+from rich.table import Table
+
+from pyintelowl import IntelOwlClientException
+from pyintelowl.cli._utils import (
+    ClickContext,
+    add_options,
+    get_success_text,
+    json_flag_option,
+)
+
+
+@click.group(help="Manage playbooks")
+def playbooks():
+    pass
+
+
+def _display_playbook(data):
+    style = "[bold #31DDCF]"
+    tags = ", ".join(
+        data["tags"]
+    )  # this is a [str], not a complex object like in job API
+    console = Console()
+    console.print(data)
+    r = Group(
+        f"{style}Playbook ID:[/] {str(data['id'])}",
+        f"{style}Name:[/] {data['name']}",
+        f"{style}Tags:[/] {tags}",
+        f"{style}TLP:[/] {data['tlp']}",
+        f"{style}Analyzers:[/] {data['analyzers']}",
+        f"{style}Connectors:[/] {data['connectors']}",
+        f"{style}Pivots:[/] {data['pivots']}",
+        f"{style}Visualizers:[/] {data['visualizers']}",
+        f"{style}Runtime configuration:[/] {data['runtime_configuration']}",
+        f"{style}For Organizations:[/] {data['for_organization']}",
+        f"{style}Disabled:[/] {data['disabled']}",
+        f"{style}Starting:[/] {data['starting']}",
+        f"{style}Description:[/] {data['description']}",
+    )
+    return Panel(r, title="Playbook attributes")
+
+
+@playbooks.command(help="Tabular print playbook attributes for a playbook name")
+@click.argument("playbook_name", type=str)
+@add_options(json_flag_option)
+@click.pass_context
+def view(
+    ctx: ClickContext,
+    playbook_name: str,
+    as_json: bool,
+):
+    ctx.obj.logger.info(f"Requesting Playbook [underline blue]{playbook_name}[/]..")
+    try:
+        ans = ctx.obj.get_playbook_by_name(playbook_name)
+    except IntelOwlClientException as e:
+        ctx.obj.logger.fatal(str(e))
+
+    if as_json:
+        rprint(json.dumps(ans, indent=4))
+    else:
+        _display_playbook(ans)
+
+
+def _display_all_playbooks(logger, rows):
+    console = Console()
+    table = Table(show_header=True, title="List of Playbooks", box=box.DOUBLE_EDGE)
+    header_style = "bold blue"
+    headers: [] = [
+        "id",
+        "name",
+        "tags",
+        "tlp",
+        "analyzers",
+        "connectors",
+        "pivots",
+        "visualizers",
+        "runtime_configuration",
+        "for_organization",
+        "disabled",
+        "starting",
+        "description",
+    ]
+    [table.add_column(header=header, header_style=header_style) for header in headers]
+    try:
+        for el in rows:
+            table.add_row(
+                str(el["id"]),
+                el["name"],
+                ", ".join([str(tag) for tag in el["tags"]]),
+                el["tlp"],
+                ", ".join([str(tag) for tag in el["analyzers"]]),
+                ", ".join([str(tag) for tag in el["connectors"]]),
+                ", ".join([str(tag) for tag in el["pivots"]]),
+                ", ".join([str(tag) for tag in el["visualizers"]]),
+                str(el["runtime_configuration"]),
+                get_success_text(el["for_organization"]),
+                get_success_text(el["disabled"]),
+                get_success_text(el["starting"]),
+                el["description"],
+            )
+        console.print(table, justify="center")
+    except Exception as e:
+        logger.fatal(e, exc_info=True)
+
+
+@playbooks.command(help="List all playbooks")
+@add_options(json_flag_option)
+@click.pass_context
+def ls(ctx: ClickContext, as_json: bool):
+    ctx.obj.logger.info("Requesting list of playbooks..")
+    try:
+        ans = ctx.obj.get_all_playbooks()
+        results = ans.get("results", [])
+        ctx.obj.logger.info(results)
+        if as_json:
+            rprint(json.dumps(results, indent=4))
+        else:
+            _display_all_playbooks(ctx.obj.logger, results)
+    except IntelOwlClientException as e:
+        ctx.obj.logger.fatal(str(e))

pyintelowl/pyintelowl.py

@@ -1181,3 +1181,48 @@ def connector_healthcheck(self, connector_name: str) -> Optional[bool]:
         url = self.instance + f"/api/connector/{connector_name}/healthcheck"
         response = self.__make_request("GET", url=url)
         return response.json().get("status", None)
+
+    def get_playbook_by_name(self, playbook_name: str) -> Dict[str, Any]:
+        """Fetch playbook info by its name.
+        Endpoint: ``/api/playbook/{playbook_name}``
+
+        Args:
+            playbook_name (str): Playbook name to retrieve
+
+        Raises:
+            IntelOwlClientException: on client/HTTP error
+
+        Returns:
+            Dict[str, Any]: JSON body.
+        """
+        url = self.instance + "/api/playbook/" + playbook_name
+        response = self.__make_request("GET", url=url)
+        return response.json()
+
+    def get_all_playbooks(self) -> Dict[str, Any]:
+        """Fetch all playbooks info.
+        Endpoint: ``/api/playbook``
+
+        Raises:
+            IntelOwlClientException: on client/HTTP error
+
+        Returns:
+            Dict[str, Any]: JSON body.
+        """
+        url = self.instance + "/api/playbook"
+        response = self.__make_request("GET", url=url)
+        return response.json()
+
+    def disable_playbook_for_org(self, playbook_name: str):
+        """Disables the plugin for the organization of the user.
+        Endpoint: ``/api/playbook/{playbook_name}/organization``
+
+        Args:
+            playbook_name (str): Playbook name to disable for org
+
+        Raises:
+            IntelOwlClientException: on client/HTTP error
+        """
+        url = self.instance + "/api/playbook/" + playbook_name + "/organization"
+        # this call doesn't have a response
+        self.__make_request("POST", url=url)

tests/__init__.py

@@ -1,5 +1,6 @@
 MOCK_CONNECTIONS = True
 TEST_JOB_ID = 1
+TEST_PLAYBOOK_NAME = "Playbook1"
 TEST_INVESTIGATION_ID = 1
 TEST_IP = "8.8.8.8"
 TEST_DOMAIN = "www.google.com"

tests/mocked_requests.py

@@ -401,3 +401,101 @@ def mocked_get_all_investigations(*args, **kwargs):
         200,
         "/api/investigation",
     )
+
+
+def mocked_get_playbook_by_name(*args, **kwargs):
+    return MockResponse(
+        {
+            "id": 1,
+            "type": ["test"],
+            "analyzers": ["Analyzer1"],
+            "connectors": [],
+            "pivots": [],
+            "visualizers": [],
+            "runtime_configuration": {
+                "pivots": {},
+                "analyzers": {},
+                "connectors": {},
+                "visualizers": {},
+            },
+            "scan_mode": 2,
+            "scan_check_time": "1:00:00:00",
+            "tags": [],
+            "tlp": "CLEAR",
+            "weight": 1,
+            "is_editable": False,
+            "for_organization": False,
+            "name": "Playbook1",
+            "description": "test",
+            "disabled": False,
+            "starting": True,
+            "owner": None,
+        },
+        200,
+        "/api/playbook/Playbook1",
+    )
+
+
+def mocked_get_all_playbooks(*args, **kwargs):
+    return MockResponse(
+        {
+            "count": 2,
+            "total_pages": 1,
+            "results": [
+                {
+                    "id": 1,
+                    "type": ["test"],
+                    "analyzers": ["Analyzer1"],
+                    "connectors": [],
+                    "pivots": [],
+                    "visualizers": [],
+                    "runtime_configuration": {
+                        "pivots": {},
+                        "analyzers": {},
+                        "connectors": {},
+                        "visualizers": {},
+                    },
+                    "scan_mode": 2,
+                    "scan_check_time": "1:00:00:00",
+                    "tags": [],
+                    "tlp": "CLEAR",
+                    "weight": 1,
+                    "is_editable": False,
+                    "for_organization": False,
+                    "name": "Playbook1",
+                    "description": "test",
+                    "disabled": False,
+                    "starting": True,
+                    "owner": None,
+                },
+                {
+                    "id": 2,
+                    "type": ["test2"],
+                    "analyzers": ["Analyzer2"],
+                    "connectors": [],
+                    "pivots": ["Pivot1"],
+                    "visualizers": ["Visualizer"],
+                    "runtime_configuration": {
+                        "pivots": {},
+                        "analyzers": {},
+                        "connectors": {},
+                        "visualizers": {},
+                    },
+                    "scan_mode": 2,
+                    "scan_check_time": "1:00:00:00",
+                    "tags": [],
+                    "tlp": "AMBER",
+                    "weight": 1,
+                    "is_editable": False,
+                    "for_organization": False,
+                    "name": "Playbook2",
+                    "description": "test",
+                    "disabled": False,
+                    "starting": True,
+                    "owner": None,
+                },
+            ],
+        },
+        200,
+        "/api/playbook",
+    )

tests/test_playbooks.py

@@ -0,0 +1,34 @@
+from unittest.mock import patch
+
+from tests.mocked_requests import mocked_get_all_playbooks, mocked_get_playbook_by_name
+from tests.utils import BaseTest, mock_connections
+
+
+class TestPlaybooks(BaseTest):
+    @mock_connections(
+        patch("requests.Session.get", side_effect=mocked_get_playbook_by_name)
+    )
+    def test_get_playbook_by_name(self, mock_requests):
+        playbook = self.client.get_playbook_by_name(self.playbook_name)
+        self.assertEqual(playbook.get("name", ""), "Playbook1")
+        self.assertEqual(playbook.get("type", []), ["test"])
+        self.assertTrue(playbook.get("analyzers", []))
+        self.assertEqual(playbook.get("analyzers", []), ["Analyzer1"])
+        self.assertFalse(playbook.get("connectors", []))
+        self.assertFalse(playbook.get("pivots", []))
+        self.assertFalse(playbook.get("visualizers", []))
+        self.assertEqual(playbook.get("description", ""), "test")
+
+    @mock_connections(
+        patch("requests.Session.get", side_effect=mocked_get_all_playbooks)
+    )
+    def test_get_all_playbooks(self, mock_requests):
+        playbooks = self.client.get_all_playbooks()
+        self.assertEqual(playbooks.get("count", 0), 2)
+        self.assertEqual(playbooks.get("total_pages", 0), 1)
+        self.assertTrue(playbooks.get("results", []))
+        self.assertEqual(len(playbooks.get("results", [])), 2)
+
+        results = playbooks.get("results", [])
+        self.assertEqual(results[0].get("name", ""), "Playbook1")
+        self.assertEqual(results[1].get("name", ""), "Playbook2")

tests/utils.py

@@ -15,6 +15,7 @@
     TEST_INVESTIGATION_ID,
     TEST_IP,
     TEST_JOB_ID,
+    TEST_PLAYBOOK_NAME,
     TEST_URL,
 )
 
@@ -56,6 +57,7 @@ class BaseTest(TestCase):
     def setUp(self):
         self.client = IntelOwl("test-token", "test-url")
         self.job_id = TEST_JOB_ID
+        self.playbook_name = TEST_PLAYBOOK_NAME
         self.investigation_id = TEST_INVESTIGATION_ID
         self.ip = TEST_IP
         self.url = TEST_URL