holmesgpt 0.11.5__py3-none-any.whl

holmes/plugins/toolsets/atlas_mongodb/mongodb_atlas.py

@@ -0,0 +1,307 @@
+import requests  # type: ignore
+import logging
+from typing import Any, Dict, Tuple, List
+from holmes.core.tools import (
+    CallablePrerequisite,
+    Tool,
+    ToolParameter,
+    Toolset,
+    ToolsetTag,
+)
+
+from pydantic import BaseModel, PrivateAttr
+from holmes.core.tools import StructuredToolResult, ToolResultStatus
+from requests.auth import HTTPDigestAuth  # type: ignore
+import gzip
+import io
+from datetime import datetime, timedelta, timezone
+import os
+from collections import Counter
+
+
+class MongoDBConfig(BaseModel):
+    public_key: str
+    private_key: str
+    project_id: str
+
+
+# https://www.mongodb.com/docs/atlas/reference/api-resources-spec/v2/
+class MongoDBAtlasToolset(Toolset):
+    name: str = "MongoDBAtlas"
+    description: str = "The MongoDB Atlas API allows access to Mongodb projects and processes. You can find logs, alerts, events, slow queries and various metrics to understand the state of Mongodb projects."
+    docs_url: str = (
+        "https://www.mongodb.com/docs/atlas/reference/api-resources-spec/v2/"
+    )
+    icon_url: str = "https://webimages.mongodb.com/_com_assets/cms/kuyjf3vea2hg34taa-horizontal_default_slate_blue.svg?auto=format%252Ccompress"
+    tags: List[ToolsetTag] = [ToolsetTag.CORE]
+    _session: requests.Session = PrivateAttr(default=requests.Session())
+
+    def __init__(self):
+        super().__init__(
+            prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)],
+            experimental=True,
+            tools=[
+                ReturnProjectAlerts(toolset=self),
+                ReturnProjectProcesses(toolset=self),
+                ReturnProjectSlowQueries(toolset=self),
+                ReturnEventsFromProject(toolset=self),
+                ReturnLogsForProcessInProject(toolset=self),
+                ReturnEventTypeFromProject(toolset=self),
+            ],
+        )
+        instructions_filepath = os.path.abspath(
+            os.path.join(os.path.dirname(__file__), "instructions.jinja2")
+        )
+        self._load_llm_instructions(jinja_template=f"file://{instructions_filepath}")
+
+    def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]:
+        if not config:
+            return False, "Missing config credentials."
+
+        try:
+            self.config: Dict = MongoDBConfig(**config).model_dump()
+            self._session.headers.update(
+                {"Accept": "application/vnd.atlas.2025-03-12+json"}
+            )
+            self._session.auth = HTTPDigestAuth(
+                self.config.get("public_key"),
+                self.config.get("private_key"),
+            )
+            return True, ""
+        except Exception:
+            logging.exception(
+                "Invalid Atlas config. Failed to set up MongoDBAtlas toolset"
+            )
+            return False, "Invalid Atlas config"
+
+    def get_example_config(self) -> Dict[str, Any]:
+        return {}
+
+
+class MongoDBAtlasBaseTool(Tool):
+    toolset: MongoDBAtlasToolset
+
+    def return_result(
+        self, response: requests.Response, params: Any, field: str = "results"
+    ) -> StructuredToolResult:
+        response.raise_for_status()
+        if response.ok:
+            res = response.json()
+            return StructuredToolResult(
+                status=ToolResultStatus.SUCCESS
+                if res.get(field, [])
+                else ToolResultStatus.NO_DATA,
+                data=res,
+                params=params,
+            )
+        else:
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Failed {self.name}.\n{response.text}",
+                return_code=response.status_code,
+                params=params,
+            )
+
+    def get_parameterized_one_liner(self, params) -> str:
+        return f"MongoDB {self.name} project {self.toolset.config.get('project_id')} {params}"
+
+
+# https://www.mongodb.com/docs/atlas/reference/api-resources-spec/v2/#tag/Alerts/operation/listAlertsByAlertConfigurationId
+class ReturnProjectAlerts(MongoDBAtlasBaseTool):
+    name: str = "atlas_return_project_alerts"
+    description: str = "Returns all project alerts. These alerts apply to all components in one project. You receive an alert when a monitored component meets or exceeds a value you set."
+
+    def _invoke(self, params: Any) -> StructuredToolResult:
+        try:
+            url = "https://cloud.mongodb.com/api/atlas/v2/groups/{project_id}/alerts".format(
+                project_id=self.toolset.config.get("project_id")
+            )
+            response = self.toolset._session.get(url=url)
+            return self.return_result(response, params)
+        except Exception as e:
+            logging.exception(self.get_parameterized_one_liner(params))
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                data=f"Exception {self.name}: {str(e)}",
+                params=params,
+            )
+
+
+# https://www.mongodb.com/docs/atlas/reference/api-resources-spec/v2/#tag/Monitoring-and-Logs/operation/listAtlasProcesses
+class ReturnProjectProcesses(MongoDBAtlasBaseTool):
+    name: str = "atlas_return_project_processes"
+    description: str = "Returns details of all processes for the specified project. Useful for getting logs and data for specific project"
+
+    def _invoke(self, params: Any) -> StructuredToolResult:
+        try:
+            url = "https://cloud.mongodb.com/api/atlas/v2/groups/{project_id}/processes".format(
+                project_id=self.toolset.config.get("project_id")
+            )
+            response = self.toolset._session.get(url)
+            return self.return_result(response, params)
+        except Exception as e:
+            logging.exception(self.get_parameterized_one_liner(params))
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Exception {self.name}: {str(e)}",
+                params=params,
+            )
+
+
+# https://www.mongodb.com/docs/atlas/reference/api-resources-spec/v2/#tag/Performance-Advisor/operation/listSlowQueries
+class ReturnProjectSlowQueries(MongoDBAtlasBaseTool):
+    name: str = "atlas_return_project_processes_slow_queries"
+    description: str = "Returns log lines for slow queries that the Performance Advisor and Query Profiler identified for a specific process in a specific project. requires fetching the project processes first. returns queries from the last 24 hours."
+    url: str = "https://cloud.mongodb.com/api/atlas/v2/groups/{project_id}/processes/{process_id}/performanceAdvisor/slowQueryLogs?includeMetrics=true"
+    parameters: Dict[str, ToolParameter] = {
+        "process_id": ToolParameter(
+            description="Combination of host and port that serves the MongoDB process. call tool atlas_return_project_processes tool to get host+port of project procecess.",
+            type="string",
+            required=True,
+        ),
+    }
+
+    def _invoke(self, params: Any) -> StructuredToolResult:
+        try:
+            url = self.url.format(
+                project_id=self.toolset.config.get("project_id"),
+                process_id=params.pop("process_id", ""),
+            )
+            response = self.toolset._session.get(url)
+            return self.return_result(response, params, "slowQueries")
+        except Exception as e:
+            logging.exception(self.get_parameterized_one_liner(params))
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Exception {self.name}: {str(e)}",
+                params=params,
+            )
+
+
+# https://www.mongodb.com/docs/atlas/reference/api-resources-spec/v2/#tag/Events/operation/listProjectEvents
+class ReturnEventsFromProject(MongoDBAtlasBaseTool):
+    name: str = "atlas_return_events_from_project"
+    description: str = "Returns all events occurrences for the specified project. Events identify significant database, security activities or status changes. can only query the last 4 hours."
+    url: str = "https://cloud.mongodb.com/api/atlas/v2/groups/{projectId}/events"
+
+    def _invoke(self, params: Any) -> StructuredToolResult:
+        params.update({"itemsPerPage": 500})
+        try:
+            now_utc = datetime.now(timezone.utc)
+            four_hours_ago = now_utc - timedelta(hours=4)
+            iso_timestamp = four_hours_ago.isoformat()
+            url = self.url.format(projectId=self.toolset.config.get("project_id"))
+            response = self.toolset._session.get(
+                url=url,
+                params={"minDate": iso_timestamp},
+            )
+            response.raise_for_status()
+            if response.ok:
+                res = response.json()
+                events_counter = Counter(
+                    [event.get("eventTypeName") for event in res.get("results", [])]
+                )
+                data = f"last 4 hours eventTypeName and # of occurrences list: {events_counter} \n to get more information about a given eventTypeName call atlas_return_events_type_from_project"
+                status = (
+                    ToolResultStatus.SUCCESS
+                    if events_counter
+                    else ToolResultStatus.NO_DATA
+                )
+                return StructuredToolResult(status=status, data=data, params=params)
+            else:
+                return StructuredToolResult(
+                    status=ToolResultStatus.ERROR,
+                    error=f"Failed {self.name}. \n{response.text}",
+                    return_code=response.status_code,
+                    params=params,
+                )
+        except Exception as e:
+            logging.exception(self.get_parameterized_one_liner(params))
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Exception {self.name}: {str(e)}",
+                params=params,
+            )
+
+
+# https://www.mongodb.com/docs/atlas/reference/api-resources-spec/v2/#tag/Monitoring-and-Logs/operation/getHostLogs
+class ReturnLogsForProcessInProject(MongoDBAtlasBaseTool):
+    name: str = "atlas_return_logs_for_host_in_project"
+    description: str = "Returns log messages for the specified host for the specified project of the last 1 hour."
+    url: str = "https://cloud.mongodb.com/api/atlas/v2/groups/{project_id}/clusters/{process_id}/logs/mongodb.gz"
+    parameters: Dict[str, ToolParameter] = {
+        "hostName": ToolParameter(
+            description="The host must be the hostname, FQDN, IPv4 address, or IPv6 address of the host that runs the MongoDB process (mongod or mongos).",
+            type="string",
+            required=True,
+        ),
+    }
+
+    def _invoke(self, params: Any) -> StructuredToolResult:
+        one_hour_ago = datetime.now(timezone.utc) - timedelta(hours=1)
+        try:
+            url = self.url.format(
+                project_id=self.toolset.config.get("project_id"),
+                process_id=params.get("hostName", ""),
+            )
+            response = self.toolset._session.get(
+                url=url,
+                headers={"Accept": "application/vnd.atlas.2025-03-12+gzip"},
+                params={"startDate": int(one_hour_ago.timestamp())},
+            )
+            response.raise_for_status()
+            if response.ok:
+                with gzip.GzipFile(fileobj=io.BytesIO(response.content)) as gz:
+                    text_data = gz.read().decode("utf-8")
+                return StructuredToolResult(
+                    status=ToolResultStatus.SUCCESS, data=text_data, params=params
+                )
+            else:
+                return StructuredToolResult(
+                    status=ToolResultStatus.ERROR,
+                    error=f"Failed {self.name}. \n{response.text}",
+                    return_code=response.status_code,
+                    params=params,
+                )
+        except Exception as e:
+            logging.exception(self.get_parameterized_one_liner(params))
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Exception {self.name}: {str(e)}",
+                params=params,
+            )
+
+
+# https://www.mongodb.com/docs/atlas/reference/api-resources-spec/v2/#tag/Events/operation/listProjectEvents
+class ReturnEventTypeFromProject(MongoDBAtlasBaseTool):
+    name: str = "atlas_return_events_type_from_project"
+    description: str = "Returns all events of specific EventType for the specified project. can only query the last 4 hours."
+    url: str = "https://cloud.mongodb.com/api/atlas/v2/groups/{projectId}/events"
+    parameters: Dict[str, ToolParameter] = {
+        "eventType": ToolParameter(
+            description="A label of an eventType, all capital letters with snake case. examples:  INSIDE_METRIC_THRESHOLD, PRIMARY_ELECTED and DATA_EXPLORER. NEVER call this before first calling atlas_return_events_from_project to get a list of last 4 hours eventTypes.",
+            type="string",
+            required=True,
+        ),
+    }
+
+    def _invoke(self, params: Any) -> StructuredToolResult:
+        try:
+            url = self.url.format(projectId=self.toolset.config.get("project_id"))
+
+            now_utc = datetime.now(timezone.utc)
+            four_hours_ago = now_utc - timedelta(hours=4)
+            iso_timestamp = four_hours_ago.isoformat()
+            params.update({"itemsPerPage": 500, "minDate": iso_timestamp})
+            response = self.toolset._session.get(
+                url=url,
+                params=params,
+            )
+            return self.return_result(response, params)
+        except Exception as e:
+            logging.exception(self.get_parameterized_one_liner(params))
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Exception {self.name}: {str(e)}",
+                params=params,
+            )

holmes/plugins/toolsets/aws.yaml

@@ -0,0 +1,76 @@
+toolsets:
+  aws/security:
+    description: "Set of tools to audit AWS security"
+    docs_url: "https://docs.robusta.dev/master/configuration/holmesgpt/toolsets/aws.html#security"
+    icon_url: "https://upload.wikimedia.org/wikipedia/commons/9/93/Amazon_Web_Services_Logo.svg"
+    tags:
+      - cli
+    prerequisites:
+      - command: "aws sts get-caller-identity"
+
+    tools:
+      - name: "aws_cloudtrail_event_lookup"
+        description: "Fetches events of a specified type from AWS CloudTrail along with the users that called them"
+        user_description: "get {{ EVENT_NAME }} events"
+        command: |
+          echo "EventName,EventId,EventTime,Username,AccessKeyId,ip,userID"
+          aws cloudtrail lookup-events \
+            --lookup-attributes AttributeKey=EventName,AttributeValue="{{ EVENT_NAME }}" \
+            --start-time $(date -v -7d -u +%Y-%m-%dT%H:%M:%SZ) \
+            --end-time $(date -u +%Y-%m-%dT%H:%M:%SZ) \
+            --query 'Events[*].{EventName:EventName,EventId:EventId,EventTime:EventTime,Username:Username,AccessKeyId:AccessKeyId,ip:CloudTrailEvent.sourceIPAddress,userID:CloudTrailEvent.userIdentity.sessionContext.sessionIssuer.userName}' \
+            --output table
+
+      - name: "aws_cloudtrail_event_details"
+        description: "Fetches and returns full event details for an AWS cloudtrail event in JSON format given an event ID"
+        user_description: "looking up event {{ EVENT_ID }}"
+        command: |
+          aws cloudtrail lookup-events \
+            --lookup-attributes AttributeKey=EventId,AttributeValue="{{ EVENT_ID }}" \
+            --query 'Events[0]' --output json
+
+      - name: "aws_user_audit_logs"
+        description: "Fetches audit logs for a specified user from AWS CloudTrail in past 24 hours. Provide username as was outputed by aws_event_lookup or aws_event_details"
+        user_description: "get audit logs for {{ UserName }}"
+        command: |
+          aws cloudtrail lookup-events \
+          --lookup-attributes AttributeKey=Username,AttributeValue="{{ UserName }}" \
+          --start-time $(date -v -1d -u +%Y-%m-%dT%H:%M:%SZ) \
+          --end-time $(date -u +%Y-%m-%dT%H:%M:%SZ) \
+          --query 'Events[*].{EventName:EventName,EventSource:EventSource,EventId:EventId,EventTime:EventTime,Username:Username,AccessKeyId:AccessKeyId,ip:CloudTrailEvent.sourceIPAddress,userID:CloudTrailEvent.userIdentity.sessionContext.sessionIssuer.userName}' \
+          --output table
+
+  aws/rds:
+    description: "Read access to Amazon RDS resources"
+    docs_url: "https://docs.robusta.dev/master/configuration/holmesgpt/toolsets/aws.html#rds"
+    icon_url: "https://upload.wikimedia.org/wikipedia/commons/9/93/Amazon_Web_Services_Logo.svg"
+    tags:
+      - cli
+    prerequisites:
+      - command: "aws sts get-caller-identity"
+
+    tools:
+      - name: "aws_rds_describe_events"
+        description: "Runs aws rds describe-events"
+        user_description: "fetch rds events"
+        command: "aws rds describe-events"
+
+      - name: "aws_rds_describe_instance"
+        description: "Get the configuration of a RDS instance"
+        user_description: "Get the configuration of a RDS instance"
+        command: "aws rds describe-db-instances --db-instance-identifier '{{ db_instance_identifier }}'"
+
+      - name: "aws_rds_describe_instances"
+        description: "Runs aws rds describe-db-instances"
+        user_description: "fetch rds instances"
+        command: "aws rds describe-db-instances"
+
+      - name: "aws_rds_describe_logs"
+        description: "Describe all available logs for an AWS RDS instance."
+        user_description: "list available RDS logs (e.g. slow query logs)"
+        command: "aws rds describe-db-log-files --db-instance-identifier '{{ db_instance_identifier }}'"
+
+      - name: "aws_rds_fetch_log_by_name"
+        description: "Fetch a specific log for an AWS RDS instance by log file name."
+        user_description: "fetch a specific RDS log"
+        command: "aws rds download-db-log-file-portion --db-instance-identifier '{{ db_instance_identifier }}' --log-file-name '{{ log_file_name }}' --starting-token 0"