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

prelude_sdk_beta/controllers/http_controller.py

@@ -0,0 +1,83 @@
+import os
+import requests
+
+from requests.adapters import HTTPAdapter, Retry
+
+
+PRELUDE_BACKOFF_FACTOR = int(os.getenv("PRELUDE_BACKOFF_FACTOR", 30))
+PRELUDE_BACKOFF_TOTAL = int(os.getenv("PRELUDE_BACKOFF_TOTAL", 0))
+
+
+class HttpController(object):
+    def __init__(self, account):
+        self._session = requests.Session()
+        self.account = account
+
+        retry = Retry(
+            total=PRELUDE_BACKOFF_TOTAL,
+            backoff_factor=PRELUDE_BACKOFF_FACTOR,
+            status_forcelist=[429],
+        )
+
+        self._session.mount("http://", HTTPAdapter(max_retries=retry))
+        self._session.mount("https://", HTTPAdapter(max_retries=retry))
+
+    def resolve_enums(self, data, enum_params: list[tuple]):
+        for [enum_class, key] in enum_params:
+            self._resolve_enum(data, enum_class, key)
+
+    def _resolve_enum(self, data, enum_class, key):
+        if isinstance(data, list):
+            for item in data:
+                if isinstance(item, dict):
+                    self._resolve_enum(item, enum_class, key)
+        elif isinstance(data, dict):
+            for k, v in data.items():
+                if k == key:
+                    if isinstance(v, list):
+                        for i, item in enumerate(v):
+                            v[i] = enum_class[item].name
+                    elif v is not None:
+                        data[k] = enum_class[v].name
+                elif isinstance(v, dict) or isinstance(v, list):
+                    self._resolve_enum(v, enum_class, key)
+
+    def get(self, url, retry=True, **kwargs):
+        res = self._session.get(url, **kwargs)
+        if res.status_code == 200:
+            return res
+        if res.status_code == 401 and retry and self.account.token_location:
+            self.account.refresh_tokens()
+            self.account.update_auth_header()
+            return self.get(url, retry=False, **kwargs)
+        raise Exception(res.text)
+
+    def post(self, url, retry=True, **kwargs):
+        res = self._session.post(url, **kwargs)
+        if res.status_code == 200:
+            return res
+        if res.status_code == 401 and retry and self.account.token_location:
+            self.account.refresh_tokens()
+            self.account.update_auth_header()
+            return self.post(url, retry=False, **kwargs)
+        raise Exception(res.text)
+
+    def delete(self, url, retry=True, **kwargs):
+        res = self._session.delete(url, **kwargs)
+        if res.status_code == 200:
+            return res
+        if res.status_code == 401 and retry and self.account.token_location:
+            self.account.refresh_tokens()
+            self.account.update_auth_header()
+            return self.delete(url, retry=False, **kwargs)
+        raise Exception(res.text)
+
+    def put(self, url, retry=True, **kwargs):
+        res = self._session.put(url, **kwargs)
+        if res.status_code == 200:
+            return res
+        if res.status_code == 401 and retry and self.account.token_location:
+            self.account.refresh_tokens()
+            self.account.update_auth_header()
+            return self.put(url, retry=False, **kwargs)
+        raise Exception(res.text)

prelude_sdk_beta/controllers/iam_controller.py

@@ -0,0 +1,285 @@
+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, Mode, Permission
+
+
+class IAMAccountController(HttpController):
+
+    def __init__(self, account):
+        super().__init__(account)
+
+    @verify_credentials
+    def get_account(self):
+        """Get account properties"""
+        res = self.get(
+            f"{self.account.hq}/iam/account", headers=self.account.headers, timeout=10
+        )
+        account = res.json()
+        if self.account.resolve_enums:
+            self.resolve_enums(
+                account, [(Mode, "mode"), (Permission, "permission"), (Control, "id")]
+            )
+        return account
+
+    @verify_credentials
+    def purge_account(self):
+        """Delete an account and all things in it"""
+        res = self.delete(
+            f"{self.account.hq}/iam/account", headers=self.account.headers, timeout=10
+        )
+        return res.json()
+
+    @verify_credentials
+    def update_account(self, mode: Mode = None, company: str = None, slug: str = None):
+        """Update properties on an account"""
+        body = dict()
+        if mode is not None:
+            body["mode"] = mode.name
+        if company is not None:
+            body["company"] = company
+        if slug is not None:
+            body["slug"] = slug
+
+        res = self.put(
+            f"{self.account.hq}/iam/account",
+            headers=self.account.headers,
+            json=body,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def attach_oidc(
+        self,
+        client_id: str,
+        client_secret: str,
+        issuer: str,
+        oidc_url: str,
+    ):
+        """Attach OIDC to an account"""
+        email_attr = "email"
+        if issuer == "azure":
+            email_attr = "upn"
+        body = dict(
+            client_id=client_id,
+            client_secret=client_secret,
+            email_attr=email_attr,
+            issuer=issuer,
+            oidc_url=oidc_url,
+        )
+
+        res = self.post(
+            f"{self.account.hq}/iam/account/oidc",
+            headers=self.account.headers,
+            json=body,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def detach_oidc(self):
+        """Detach OIDC to an account"""
+        res = self.delete(
+            f"{self.account.hq}/iam/account/oidc",
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def invite_user(
+        self,
+        email: str,
+        oidc: str | None,
+        permission: Permission,
+        name: str | None = None,
+    ):
+        """Invite a new user to the account"""
+        body = dict(permission=permission.name, handle=email, oidc=oidc)
+        if name:
+            body["name"] = name
+
+        res = self.post(
+            url=f"{self.account.hq}/iam/account/user",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        user = res.json()
+        if self.account.resolve_enums:
+            self.resolve_enums(user, [(Permission, "permission")])
+        return user
+
+    @verify_credentials
+    def create_service_user(self, name: str):
+        """Create a service user"""
+        body = dict(name=name)
+
+        res = self.post(
+            f"{self.account.hq}/iam/account/service_user",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def delete_service_user(self, handle: str):
+        """Delete service user"""
+        body = dict(handle=handle)
+
+        res = self.delete(
+            f"{self.account.hq}/iam/account/service_user",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def update_account_user(
+        self,
+        email: str,
+        oidc: str | None,
+        permission: Permission = None,
+    ):
+        """Update properties on an account user"""
+        body = dict(handle=email, oidc=oidc)
+        if permission is not None:
+            body["permission"] = permission.name
+
+        res = self.put(
+            f"{self.account.hq}/iam/account/user",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def remove_user(self, email: str, oidc: str | None):
+        """Remove user from the account"""
+        params = dict(handle=email, oidc=oidc)
+
+        res = self.delete(
+            f"{self.account.hq}/iam/account/user",
+            params=params,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def audit_logs(self, days: int = 7, limit: int = 1000):
+        """Get audit logs from the last X days"""
+        params = dict(days=days, limit=limit)
+        res = self.get(
+            f"{self.account.hq}/iam/audit",
+            headers=self.account.headers,
+            params=params,
+            timeout=30,
+        )
+        return res.json()
+
+    def sign_up(self, company, email, name):
+        """(NOT AVAIABLE IN PRODUCTION) Create a new user and account"""
+        body = dict(company=company, email=email, name=name)
+
+        res = self._session.post(
+            f"{self.account.hq}/iam/new_user_and_account",
+            headers=self.account.headers,
+            json=body,
+            timeout=20,
+        )
+        data = res.json()
+        if self.account.profile:
+            self.account.keychain.configure_keychain(
+                account=data["account_id"],
+                handle=data["user_id"],
+                hq=self.account.hq,
+                profile=self.account.profile,
+            )
+        return data
+
+
+class IAMUserController(HttpController):
+
+    def __init__(self, account):
+        super().__init__(account)
+
+    @verify_credentials
+    def list_accounts(self):
+        """List all accounts for your user"""
+        res = self.get(
+            f"{self.account.hq}/iam/user/account",
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def purge_user(self):
+        """Delete your user"""
+        res = self.delete(
+            f"{self.account.hq}/iam/user", headers=self.account.headers, timeout=10
+        )
+        return res.json()
+
+    @verify_credentials
+    def update_user(
+        self,
+        name: str = None,
+    ):
+        """Update properties on a user"""
+        body = dict()
+        if name is not None:
+            body["name"] = name
+
+        res = self.put(
+            f"{self.account.hq}/iam/user",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    def forgot_password(self):
+        """Send a forgot password email"""
+        body = dict(handle=self.account.handle)
+
+        res = self.post(
+            f"{self.account.hq}/iam/user/forgot_password",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    def confirm_forgot_password(self, confirmation_code: str, new_password: str):
+        """Change a password using confirmation code"""
+        body = dict(
+            handle=self.account.handle,
+            confirmation_code=confirmation_code,
+            password=new_password,
+        )
+
+        res = self.post(
+            f"{self.account.hq}/iam/user/forgot_password",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def change_password(self, current_password: str, new_password: str):
+        """Change your password"""
+        body = dict(current_password=current_password, new_password=new_password)
+
+        res = self.post(
+            f"{self.account.hq}/iam/user/change_password",
+            json=body,
+            headers=self.account.headers,
+            timeout=10,
+        )
+        return res.json()

prelude_sdk_beta/controllers/jobs_controller.py

@@ -0,0 +1,37 @@
+from itertools import chain
+
+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 BackgroundJobTypes, Control
+
+
+class JobsController(HttpController):
+
+    def __init__(self, account):
+        super().__init__(account)
+
+    @verify_credentials
+    def job_statuses(self):
+        """Get job statuses"""
+        res = self.get(
+            f"{self.account.hq}/jobs/statuses", headers=self.account.headers, timeout=30
+        )
+        jobs = res.json()
+        if self.account.resolve_enums:
+            self.resolve_enums(jobs, [(Control, "control")])
+        return jobs
+
+    @verify_credentials
+    def job_status(self, job_id: str):
+        """Get job status given job ID"""
+        res = self.get(
+            f"{self.account.hq}/jobs/statuses/{job_id}",
+            headers=self.account.headers,
+            timeout=30,
+        )
+        job = res.json()
+        if self.account.resolve_enums:
+            self.resolve_enums(
+                job, [(Control, "control"), (BackgroundJobTypes, "job_type")]
+            )
+        return job

prelude_sdk_beta/controllers/partner_controller.py

@@ -0,0 +1,154 @@
+from datetime import datetime, timezone
+
+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
+
+
+class PartnerController(HttpController):
+
+    def __init__(self, account):
+        super().__init__(account)
+
+    @verify_credentials
+    def attach(
+        self,
+        partner: Control,
+        api: str,
+        user: str,
+        secret: str,
+        name: str | None = None,
+        instance_id: str | None = None,
+    ):
+        """Attach a partner to your account"""
+        params = dict()
+        if name:
+            params["name"] = name
+        if api:
+            params["api"] = api
+        if user:
+            params["user"] = user
+        if secret:
+            params["secret"] = secret
+        extra = f"/{instance_id}" if instance_id else ""
+        res = self.post(
+            f"{self.account.hq}/partner/{partner.name}{extra}",
+            headers=self.account.headers,
+            json=params,
+            timeout=10,
+        )
+        return res.json()
+
+    @verify_credentials
+    def detach(self, partner: Control, instance_id: str):
+        """Detach a partner from your Detect account"""
+        res = self.delete(
+            f"{self.account.hq}/partner/{partner.name}/{instance_id}",
+            headers=self.account.headers,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def block(self, partner: Control, test_id: str):
+        """Report to a partner to block a test"""
+        params = dict(test_id=test_id)
+        res = self.post(
+            f"{self.account.hq}/partner/block/{partner.name}",
+            headers=self.account.headers,
+            json=params,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def endpoints(
+        self,
+        partner: Control,
+        platform: str,
+        hostname: str = "",
+        offset: int = 0,
+        count: int = 100,
+    ):
+        """Get a list of endpoints from a partner"""
+        params = dict(platform=platform, hostname=hostname, offset=offset, count=count)
+        res = self.get(
+            f"{self.account.hq}/partner/endpoints/{partner.name}",
+            headers=self.account.headers,
+            params=params,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def deploy(self, partner: Control, host_ids: list):
+        """Deploy probes on all specified partner endpoints"""
+        params = dict(host_ids=host_ids)
+        res = self.post(
+            f"{self.account.hq}/partner/deploy/{partner.name}",
+            headers=self.account.headers,
+            json=params,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def list_reports(self, partner: Control, test_id: str | None):
+        """Get reports to a partner for a test"""
+        params = dict(test_id=test_id) if test_id else dict()
+        res = self.get(
+            f"{self.account.hq}/partner/reports/{partner.name}",
+            headers=self.account.headers,
+            json=params,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def observed_detected(self, test_id: str | None = None, hours: int | None = None):
+        """Get observed_detected stats"""
+        params = dict()
+        if test_id:
+            params["test_id"] = test_id
+        if hours:
+            params["start_epoch_ms"] = (
+                datetime.now(timezone.utc).timestamp() - hours * 60 * 60
+            ) * 1000
+
+        res = self.get(
+            f"{self.account.hq}/partner/observed_detected",
+            headers=self.account.headers,
+            json=params,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def list_advisories(
+        self, partner: Control, start: str = None, limit: int = None, offset: int = None
+    ):
+        """Get advisory reports provided by a partner"""
+        params = dict()
+        if start:
+            params["start"] = start
+        if limit:
+            params["limit"] = limit
+        if offset:
+            params["offset"] = offset
+        res = self.get(
+            f"{self.account.hq}/partner/advisories/{partner.name}",
+            headers=self.account.headers,
+            params=params,
+            timeout=30,
+        )
+        return res.json()
+
+    @verify_credentials
+    def partner_groups(self, partner: Control, instance_id: str):
+        """Get a list of partner groups"""
+        res = self.get(
+            f"{self.account.hq}/partner/groups/{partner.name}/{instance_id}",
+            headers=self.account.headers,
+            timeout=30,
+        )
+        return res.json()

prelude_sdk_beta/controllers/probe_controller.py

@@ -0,0 +1,14 @@
+from prelude_sdk_beta.controllers.http_controller import HttpController
+
+
+class ProbeController(HttpController):
+
+    def __init__(self, account):
+        super().__init__(account)
+
+    def download(self, name: str, dos: str):
+        """Download a probe executable"""
+        res = self.get(
+            f"{self.account.hq}/download/{name}", headers=dict(dos=dos), timeout=10
+        )
+        return res.text