enable_disabled_privilege_zone_rules.py

#!/usr/bin/env python3
"""
Enable disabled BloodHound Enterprise / BloodHound CE Privilege Zone Rules.

This script enumerates asset group tags, finds disabled tag selectors, and enables them.
By default it runs in dry-run mode. Pass --apply to make changes.

Authentication:
  1. API token signed auth:
       --token-id "$BHE_TOKEN_ID" --token-key "$BHE_TOKEN_KEY"

  2. Temporary browser JWT:
       --jwt "$BHE_JWT"

Example:
  python enable_disabled_privilege_zone_rules.py \
    --tenant-url https://example.bloodhoundenterprise.io \
    --token-id "$BHE_TOKEN_ID" \
    --token-key "$BHE_TOKEN_KEY" \
    --apply
"""

import argparse
import base64
import datetime as dt
import hashlib
import hmac
import json
import sys

import requests


def sign_headers(token_id, token_key, method, uri, body=b""):
    now = dt.datetime.now(dt.timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
    key = token_key.encode()

    op_key = hmac.new(key, f"{method.upper()}{uri}".encode(), hashlib.sha256).digest()
    date_key = hmac.new(op_key, now[:13].encode(), hashlib.sha256).digest()
    sig = hmac.new(date_key, body or b"", hashlib.sha256).digest()

    return {
        "Authorization": f"bhesignature {token_id}",
        "RequestDate": now,
        "Signature": base64.b64encode(sig).decode(),
    }


class BHE:
    def __init__(self, base_url, jwt=None, token_id=None, token_key=None):
        self.base_url = base_url.rstrip("/")
        self.jwt = jwt
        self.token_id = token_id
        self.token_key = token_key
        self.session = requests.Session()

    def request(self, method, path, params=None, payload=None):
        body = json.dumps(payload).encode() if payload is not None else None
        headers = {"Accept": "application/json"}

        if payload is not None:
            headers["Content-Type"] = "application/json"

        req = requests.Request(
            method=method,
            url=self.base_url + path,
            params=params,
            data=body,
            headers=headers,
        )
        prepped = self.session.prepare_request(req)

        # Signed auth must include the query string in the request URI.
        uri_to_sign = prepped.path_url

        if self.jwt:
            prepped.headers["Authorization"] = f"Bearer {self.jwt}"
        else:
            prepped.headers.update(
                sign_headers(
                    self.token_id,
                    self.token_key,
                    method,
                    uri_to_sign,
                    body or b"",
                )
            )

        resp = self.session.send(prepped, timeout=60)

        if not resp.ok:
            print(f"\nHTTP {resp.status_code} for {method} {uri_to_sign}", file=sys.stderr)
            print(resp.text, file=sys.stderr)

        resp.raise_for_status()
        return resp.json() if resp.content else {}

    def get_all(self, path, key, params=None, limit=1000):
        skip = 0
        while True:
            q = dict(params or {})
            q.update({"skip": skip, "limit": limit})

            data = self.request("GET", path, params=q).get("data", {})
            items = data.get(key, [])

            if not items:
                break

            yield from items

            if len(items) < limit:
                break

            skip += limit


def main():
    parser = argparse.ArgumentParser(
        description="Enable all disabled BloodHound Enterprise Privilege Zone Rules."
    )
    parser.add_argument(
        "--tenant-url",
        required=True,
        help="Example: https://shinybits.bloodhoundenterprise.io",
    )
    parser.add_argument("--jwt", help="Temporary browser JWT bearer token")
    parser.add_argument("--token-id", help="BHE API token ID")
    parser.add_argument("--token-key", help="BHE API token key")
    parser.add_argument(
        "--apply",
        action="store_true",
        help="Actually enable rules. Default is dry-run.",
    )

    args = parser.parse_args()

    if not args.jwt and not (args.token_id and args.token_key):
        parser.error("Provide either --jwt or both --token-id and --token-key")

    bhe = BHE(
        base_url=args.tenant_url,
        jwt=args.jwt,
        token_id=args.token_id,
        token_key=args.token_key,
    )

    disabled_rules = []

    tags = list(
        bhe.get_all(
            "/api/v2/asset-group-tags",
            "tags",
            params={"counts": "true"},
        )
    )

    for tag in tags:
        tag_id = tag["id"]
        tag_name = tag.get("name", "<unnamed tag>")

        selectors = bhe.get_all(
            f"/api/v2/asset-group-tags/{tag_id}/selectors",
            "selectors",
        )

        for selector in selectors:
            if selector.get("disabled_at"):
                disabled_rules.append((tag, selector))
                print(
                    f"Found disabled rule: "
                    f"tag='{tag_name}' tag_id={tag_id} "
                    f"selector='{selector.get('name')}' selector_id={selector['id']}"
                )

    print(f"\nFound {len(disabled_rules)} disabled rule(s).")

    if not args.apply:
        print("Dry-run only. Re-run with --apply to enable these rules.")
        return

    for tag, selector in disabled_rules:
        tag_id = tag["id"]
        selector_id = selector["id"]

        bhe.request(
            "PATCH",
            f"/api/v2/asset-group-tags/{tag_id}/selectors/{selector_id}",
            payload={"disabled": False},
        )

        print(f"Enabled selector_id={selector_id} on tag_id={tag_id}")

    print("\nDone.")


if __name__ == "__main__":
    main()