prelude-sdk-beta 1406__py3-none-any.whl

prelude_sdk_beta/controllers/scm_controller.py

@@ -0,0 +1,424 @@
+from prelude_sdk_beta.controllers.http_controller import HttpController
+from prelude_sdk_beta.models.account import verify_credentials
+from prelude_sdk_beta.models.codes import Control, ControlCategory, PartnerEvents, RunCode
+
+
+class ScmController(HttpController):
+    default = -1
+
+    def __init__(self, account):
+        super().__init__(account)
+
+    @verify_credentials
+    def endpoints(self, filter: str = None, orderby: str = None, top: int = None):
+        """List endpoints with SCM analysis"""
+        params = {"$filter": filter, "$orderby": orderby, "$top": top}
+        res = self.get(
+            f"{self.account.hq}/scm/endpoints",
+            headers=self.account.headers,
+            params=params,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def inboxes(self, filter: str = None, orderby: str = None, top: int = None):
+        """List inboxes with SCM analysis"""
+        params = {"$filter": filter, "$orderby": orderby, "$top": top}
+        res = self.get(
+            f"{self.account.hq}/scm/inboxes",
+            headers=self.account.headers,
+            params=params,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def users(self, filter: str = None, orderby: str = None, top: int = None):
+        """List users with SCM analysis"""
+        params = {"$filter": filter, "$orderby": orderby, "$top": top}
+        res = self.get(
+            f"{self.account.hq}/scm/users",
+            headers=self.account.headers,
+            params=params,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def technique_summary(self, techniques: str):
+        """Get policy evaluation summary by technique"""
+        res = self.get(
+            f"{self.account.hq}/scm/technique_summary",
+            params=dict(techniques=techniques),
+            headers=self.account.headers,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def evaluation_summary(
+        self,
+        endpoint_filter: str = None,
+        inbox_filter: str = None,
+        user_filter: str = None,
+        techniques: str = None,
+    ):
+        """Get policy evaluation summary"""
+        params = dict(
+            endpoints_filter=endpoint_filter,
+            inboxes_filter=inbox_filter,
+            users_filter=user_filter,
+        )
+        if techniques:
+            params["techniques"] = techniques
+        res = self.get(
+            f"{self.account.hq}/scm/evaluation_summary",
+            params=params,
+            headers=self.account.headers,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def evaluation(
+        self,
+        partner: Control,
+        instance_id: str,
+        filter: str = None,
+        techniques: str = None,
+    ):
+        """Get policy evaluations for given partner"""
+        params = {"$filter": filter}
+        if techniques:
+            params["techniques"] = techniques
+        res = self.get(
+            f"{self.account.hq}/scm/evaluations/{partner.name}/{instance_id}",
+            params=params,
+            headers=self.account.headers,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def update_evaluation(self, partner: Control, instance_id: str):
+        """Update policy evaluations for given partner"""
+        res = self.post(
+            f"{self.account.hq}/scm/evaluations/{partner.name}/{instance_id}",
+            headers=self.account.headers,
+            timeout=60,
+        )
+        return res.json()
+
+    @verify_credentials
+    def list_partner_groups(self, filter: str = None, orderby: str = None):
+        """List groups"""
+        params = {"$filter": filter, "$orderby": orderby}
+        res = self.get(
+            f"{self.account.hq}/scm/groups",
+            headers=self.account.headers,
+            params=params,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def update_partner_groups(
+        self, partner: Control, instance_id: str, group_ids: list[str]
+    ):
+        """Update groups"""
+        body = dict(group_ids=group_ids)
+        res = self.post(
+            f"{self.account.hq}/scm/groups/{partner.name}/{instance_id}",
+            headers=self.account.headers,
+            json=body,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def list_object_exceptions(self):
+        """List object exceptions"""
+        res = self.get(
+            f"{self.account.hq}/scm/exceptions/objects",
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def create_object_exception(
+        self, category: ControlCategory, filter, name=None, expires: str = None
+    ):
+        """Create an object exception"""
+        body = dict(category=category.name, filter=filter)
+        if name:
+            body["name"] = name
+        if expires:
+            body["expires"] = expires
+        res = self.post(
+            f"{self.account.hq}/scm/exceptions/objects",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def update_object_exception(
+        self, exception_id, expires=default, filter=None, name=None
+    ):
+        """Update an object exception"""
+        body = dict()
+        if expires != self.default:
+            body["expires"] = expires
+        if filter:
+            body["filter"] = filter
+        if name:
+            body["name"] = name
+        res = self.post(
+            f"{self.account.hq}/scm/exceptions/objects/{exception_id}",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def delete_object_exception(self, exception_id):
+        """Delete an object exception"""
+        res = self.delete(
+            f"{self.account.hq}/scm/exceptions/objects/{exception_id}",
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def list_policy_exceptions(self):
+        """List policy exceptions"""
+        res = self.get(
+            f"{self.account.hq}/scm/exceptions/policies",
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def put_policy_exceptions(
+        self, partner: Control, expires, instance_id: str, policy_id, setting_names
+    ):
+        """Put policy exceptions"""
+        body = dict(
+            control=partner.name,
+            expires=expires,
+            instance_id=instance_id,
+            policy_id=policy_id,
+            setting_names=setting_names,
+        )
+        res = self.put(
+            f"{self.account.hq}/scm/exceptions/policies",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def list_views(self):
+        """List views"""
+        res = self.get(
+            f"{self.account.hq}/scm/views",
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def create_view(self, category: ControlCategory, filter: str, name: str):
+        """Create a view"""
+        body = dict(category=category.name, filter=filter, name=name)
+        res = self.post(
+            f"{self.account.hq}/scm/views",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def update_view(
+        self, view_id, category: ControlCategory = None, filter=None, name=None
+    ):
+        """Update a view"""
+        body = dict()
+        if category:
+            body["category"] = category.name
+        if filter:
+            body["filter"] = filter
+        if name:
+            body["name"] = name
+        res = self.post(
+            f"{self.account.hq}/scm/views/{view_id}",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def delete_view(self, view_id):
+        """Delete a view"""
+        res = self.delete(
+            f"{self.account.hq}/scm/views/{view_id}",
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def create_threat(
+        self,
+        name,
+        description=None,
+        id=None,
+        generated=None,
+        published=None,
+        source=None,
+        source_id=None,
+        techniques=None,
+    ):
+        """Create an scm threat"""
+        body = dict(name=name)
+        if description:
+            body["description"] = description
+        if id:
+            body["id"] = id
+        if generated:
+            body["generated"] = generated
+        if published:
+            body["published"] = published
+        if source:
+            body["source"] = source
+        if source_id:
+            body["source_id"] = source_id
+        if techniques:
+            body["techniques"] = techniques
+
+        res = self.post(
+            f"{self.account.hq}/scm/threats",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def delete_threat(self, id):
+        """Delete an existing scm threat"""
+        res = self.delete(
+            f"{self.account.hq}/scm/threats/{id}",
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def get_threat(self, id):
+        """Get specific scm threat"""
+        res = self.get(
+            f"{self.account.hq}/scm/threats/{id}",
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def list_threats(self):
+        """List all scm threats"""
+        res = self.get(
+            f"{self.account.hq}/scm/threats", headers=self.account.headers, timeout=10
+        )
+        return res.json()
+
+    @verify_credentials
+    def parse_threat_intel(self, file: str):
+        with open(file, "rb") as f:
+            body = f.read()
+        res = self.post(
+            f"{self.account.hq}/scm/threat-intel",
+            data=body,
+            headers=self.account.headers | {"Content-Type": "application/pdf"},
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def parse_from_partner_advisory(self, partner: Control, advisory_id: str):
+        params = dict(advisory_id=advisory_id)
+        res = self.post(
+            f"{self.account.hq}/scm/partner-advisories/{partner.name}",
+            headers=self.account.headers,
+            json=params,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def list_notifications(self):
+        res = self.get(
+            f"{self.account.hq}/scm/notifications",
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def delete_notification(self, notification_id: str):
+        res = self.delete(
+            f"{self.account.hq}/scm/notifications/{notification_id}",
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def upsert_notification(
+        self,
+        control_category: ControlCategory,
+        event: PartnerEvents,
+        run_code: RunCode,
+        scheduled_hour: int,
+        emails: list[str] = None,
+        filter: str = None,
+        id: str = None,
+        message: str = "",
+        slack_urls: list[str] = None,
+        suppress_empty: bool = True,
+        teams_urls: list[str] = None,
+        title: str = "SCM Notification",
+    ):
+        body = dict(
+            control_category=control_category.name,
+            event=event.name,
+            run_code=run_code.name,
+            scheduled_hour=scheduled_hour,
+            suppress_empty=suppress_empty,
+        )
+        if id:
+            body["id"] = id
+        if filter:
+            body["filter"] = filter
+        if emails:
+            body["email"] = dict(emails=emails, message=message, subject=title)
+        if slack_urls:
+            body["slack"] = dict(hook_urls=slack_urls, message=message)
+        if teams_urls:
+            body["teams"] = dict(hook_urls=teams_urls, message=message)
+        res = self.put(
+            f"{self.account.hq}/scm/notifications",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()

prelude_sdk_beta/models/account.py

@@ -0,0 +1,264 @@
+import configparser
+import json
+import os
+from functools import wraps
+from pathlib import Path
+
+import requests
+
+
+class Keychain:
+
+    def __init__(
+        self,
+        keychain_location: str | None = os.path.join(
+            Path.home(), ".prelude", "keychain.ini"
+        ),
+    ):
+        self.keychain_location = keychain_location
+        if self.keychain_location and not os.path.exists(self.keychain_location):
+            head, _ = os.path.split(Path(self.keychain_location))
+            Path(head).mkdir(parents=True, exist_ok=True)
+            open(self.keychain_location, "x").close()
+            self.configure_keychain("", "")
+
+    def read_keychain(self):
+        cfg = configparser.ConfigParser()
+        cfg.read(self.keychain_location)
+        return cfg
+
+    def configure_keychain(
+        self,
+        account,
+        handle,
+        hq="https://api.us1.preludesecurity.com",
+        oidc=None,
+        profile="default",
+        slug=None,
+    ):
+        cfg = self.read_keychain()
+        cfg[profile] = {"account": account, "handle": handle, "hq": hq}
+        if oidc:
+            cfg[profile]["oidc"] = oidc
+        if slug:
+            cfg[profile]["slug"] = slug
+        with open(self.keychain_location, "w") as f:
+            cfg.write(f)
+
+    def get_profile(self, profile="default") -> dict:
+        try:
+            cfg = self.read_keychain()
+            profile = next(s for s in cfg.sections() if s == profile)
+            return dict(cfg[profile].items())
+        except StopIteration:
+            raise Exception(
+                "Could not find profile %s for account in %s"
+                % (profile, self.keychain_location)
+            )
+
+
+def exchange_token(
+    account: str, handle: str, hq: str, auth_flow: str, auth_params: dict
+):
+    """
+    Two token exchange auth flows:
+    1) Password auth: auth_flow = "password", auth_params = {"password": "your_password"}
+    2) Refresh token auth: auth_flow = "refresh", auth_params = {"refresh_token": "your_refresh_token"}
+    3) Exchange an OIDC authorization code for tokens:
+       auth_flow = "oauth_code", auth_params = {"code": "your_authorization_code", "verifier": "your_verifier", "source": "cli"}
+    """
+    res = requests.post(
+        f"{hq}/iam/token",
+        headers=dict(account=account, _product="py-sdk"),
+        json=dict(auth_flow=auth_flow, handle=handle, **auth_params),
+        timeout=10,
+    )
+    if res.status_code == 401:
+        raise Exception("Error logging in: Unauthorized")
+    if not res.ok:
+        raise Exception("Error logging in: %s" % res.text)
+    return res.json()
+
+
+class Account:
+
+    @staticmethod
+    def from_keychain(profile: str = "default"):
+        """
+        Create an account object from a pre-configured profile in your keychain file
+        """
+        keychain = Keychain()
+        profile_items = keychain.get_profile(profile)
+        if any([item not in profile_items for item in ["account", "handle", "hq"]]):
+            raise ValueError(
+                "Please make sure you are using an up-to-date profile with the following fields: account, handle, hq"
+            )
+        return _Account(
+            account=profile_items["account"],
+            handle=profile_items["handle"],
+            hq=profile_items["hq"],
+            oidc=profile_items.get("oidc"),
+            profile=profile,
+            slug=profile_items.get("slug"),
+        )
+
+    @staticmethod
+    def from_token(
+        account: str,
+        handle: str,
+        token: str | None = None,
+        refresh_token: str | None = None,
+        hq: str = "https://api.us1.preludesecurity.com",
+        oidc: str | None = None,
+        slug: str | None = None,
+    ):
+        """
+        Create an account object from an access token or a refresh token
+        """
+        if not any([token, refresh_token]):
+            raise ValueError("Please provide either an access token or a refresh token")
+        if refresh_token:
+            res = exchange_token(
+                account, handle, hq, "refresh", dict(refresh_token=refresh_token)
+            )
+            token = res["token"]
+        return _Account(
+            account,
+            handle,
+            hq,
+            keychain_location=None,
+            oidc=oidc,
+            slug=slug,
+            token=token,
+            token_location=None,
+        )
+
+
+class _Account:
+
+    def __init__(
+        self,
+        account: str,
+        handle: str,
+        hq: str,
+        oidc: str | None = None,
+        profile: str | None = None,
+        slug: str | None = None,
+        token: str | None = None,
+        keychain_location: str | None = os.path.join(
+            Path.home(), ".prelude", "keychain.ini"
+        ),
+        token_location: str | None = os.path.join(
+            Path.home(), ".prelude", "tokens.json"
+        ),
+    ):
+        if token is None and token_location is None:
+            raise ValueError(
+                "Please provide either an access token or a token location"
+            )
+
+        super().__init__()
+        self.account = account
+        self.handle = handle
+        self.headers = dict(account=account, _product="py-sdk")
+        self.hq = hq
+        self.keychain = Keychain(keychain_location)
+        self.oidc = oidc
+        self.profile = profile
+        self.slug = slug
+        self.token = token
+        self.token_location = token_location
+        if self.token_location and not os.path.exists(self.token_location):
+            head, _ = os.path.split(Path(self.token_location))
+            Path(head).mkdir(parents=True, exist_ok=True)
+            with open(self.token_location, "x") as f:
+                json.dump({}, f)
+        self.source = "cli" if self.oidc else "main"
+
+    @property
+    def token_key(self):
+        return f"{self.handle}/{self.oidc}" if self.oidc else self.handle
+
+    def _read_tokens(self):
+        with open(self.token_location, "r") as f:
+            return json.load(f)
+
+    def save_new_token(self, new_tokens):
+        existing_tokens = self._read_tokens()
+        if self.token_key not in existing_tokens:
+            existing_tokens[self.token_key] = dict()
+        existing_tokens[self.token_key][self.hq] = new_tokens
+        with open(self.token_location, "w") as f:
+            json.dump(existing_tokens, f)
+
+    def _verify(self):
+        if not self.token_location:
+            raise ValueError("Please provide a token location to continue")
+        if self.profile and not any([self.handle, self.account]):
+            raise ValueError(
+                "Please configure your %s profile to continue" % self.profile
+            )
+
+    def password_login(self, password, new_password=None):
+        self._verify()
+        tokens = exchange_token(
+            self.account,
+            self.handle,
+            self.hq,
+            "password_change" if new_password else "password",
+            dict(password=password, new_password=new_password),
+        )
+        self.save_new_token(tokens)
+        return tokens
+
+    def refresh_tokens(self):
+        self._verify()
+        existing_tokens = self._read_tokens().get(self.token_key, {}).get(self.hq, {})
+        if not (refresh_token := existing_tokens.get("refresh_token")):
+            raise Exception("No refresh token found, please login first to continue")
+        tokens = exchange_token(
+            self.account,
+            self.handle,
+            self.hq,
+            "refresh",
+            dict(refresh_token=refresh_token, source=self.source),
+        )
+        tokens = existing_tokens | tokens
+        self.save_new_token(tokens)
+        return tokens
+
+    def exchange_authorization_code(self, authorization_code: str, verifier: str):
+        self._verify()
+        tokens = exchange_token(
+            self.account,
+            self.handle,
+            self.hq,
+            "oauth_code",
+            dict(code=authorization_code, verifier=verifier, source=self.source),
+        )
+        existing_tokens = self._read_tokens().get(self.token_key, {}).get(self.hq, {})
+        tokens = existing_tokens | tokens
+        self.save_new_token(tokens)
+        return tokens
+
+    def get_token(self):
+        if self.token:
+            return self.token
+
+        tokens = self._read_tokens().get(self.token_key, {}).get(self.hq, {})
+        if "token" not in tokens:
+            raise Exception("Please login to continue")
+        return tokens["token"]
+
+    def update_auth_header(self):
+        self.headers |= dict(authorization=f"Bearer {self.get_token()}")
+
+
+def verify_credentials(func):
+    @wraps(verify_credentials)
+    def handler(*args, **kwargs):
+        args[0].account.update_auth_header()
+        return func(*args, **kwargs)
+
+    handler.__wrapped__ = func
+    return handler