Spaces:

Wardline
/

cloak-api

Paused

cloak-api / copilot /platform /context.py

Upload 46 files

91b8be2 verified 2 days ago

12.5 kB

"""PlatformContext: what the tenant looks like right now.
The Node proxy builds this from the database on every request (groups, rules, block lists,
connectors, active policies) and the copilot never reads the database itself. Loading is
deliberately forgiving: a malformed entry is dropped, never raised, because a bad context
must degrade answers, not crash the service.
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Any

from . import catalog

CONNECTOR_CHANNELS = ('email', 'slack', 'teams', 'drive', 'chat', 'm365files', 's3', 'gcs', 'azure', 'github')


@dataclass
class Group:
    id: str
    name: str


@dataclass
class Rule:
    id: str
    name: str
    entities: list[str]
    keywords: list[str]
    fileTypes: list[str]
    severity: str
    builtin: bool


@dataclass
class BlockList:
    id: str
    name: str
    kind: str          # 'web' | 'app'
    entries: list[str]
    builtin: bool


@dataclass
class WebCategory:
    name: str          # as the admin named it in Web Filtering; policies store this name
    domains: int       # how many sites it holds (0 until it has been filled)


@dataclass
class Connector:
    id: str
    name: str
    hasAdminToken: bool = False        # slack: needed for the delete action
    remediationEnabled: bool = False   # email: needed for the remediate action
    provider: str = ''                 # email: 'google' | 'm365'


@dataclass
class Detector:
    """An admin-made regex detector (Policies -> Detectors). Rules reference it by label."""
    label: str         # CUSTOMER_ID: what a detection rule stores in `entities`
    name: str          # "Customer ID": what the admin calls it
    enabled: bool


@dataclass
class PolicySummary:
    id: str
    name: str
    type: str
    groupId: str
    action: str
    connectorIds: list[str]
    provider: str = ''     # email policies: 'google' | 'm365' (one saved before the split is Gmail)
    active: bool = True
    # Device and browser switches that are on ("agentControl.blockPrint"), for unblock requests.
    controls: set[str] = field(default_factory=set)
    # Sites and apps it blocks by name, and the block lists it uses: for "unblock chatgpt".
    blockedDomains: set[str] = field(default_factory=set)
    blockedApps: set[str] = field(default_factory=set)
    webLists: list[str] = field(default_factory=list)
    appLists: list[str] = field(default_factory=list)


@dataclass
class PlatformContext:
    groups: list[Group] = field(default_factory=list)
    rules: list[Rule] = field(default_factory=list)
    blockLists: list[BlockList] = field(default_factory=list)
    connectors: dict[str, list[Connector]] = field(default_factory=dict)
    # Active policies only. The copilot uses them to explain how a new one will combine.
    policies: list[PolicySummary] = field(default_factory=list)
    # Every policy, on or off: what a request may add to instead of creating another one.
    allPolicies: list[PolicySummary] = field(default_factory=list)
    # Names of every policy, active or not, so a new one never repeats a name.
    policyNames: set[str] = field(default_factory=set)
    # True when the caller told us which connectors exist. When it did not, we cannot
    # say "no connector", so connector checks are skipped rather than guessed.
    connectorsKnown: bool = False
    policiesKnown: bool = False
    # Web Filtering categories (tenant-defined). Unknown when the caller sent none.
    webCategories: list[WebCategory] = field(default_factory=list)
    webCategoriesKnown: bool = False
    # True when the caller sent each policy's switches (PolicySummary.controls).
    controlsKnown: bool = False
    # True when the caller sent what each policy blocks (sites, apps, lists).
    blocksKnown: bool = False
    # The tenant's own regex detectors, and the built-in ones an admin switched off.
    detectors: list[Detector] = field(default_factory=list)
    disabledDetectors: set[str] = field(default_factory=set)

    def entity_keys(self) -> frozenset[str]:
        """Every key a rule may hold for this tenant: the platform's plus its custom detectors."""
        if not self.detectors:
            return catalog.entity_keys()
        return catalog.entity_keys() | {x.label for x in self.detectors}

    def entity_label(self, key: str) -> str:
        own = next((x for x in self.detectors if x.label == key), None)
        return own.name if own else catalog.entity_label(key)

    def detector_off(self, key: str) -> bool:
        """Switched off on the Detectors page: a rule may list it, but it finds nothing."""
        own = next((x for x in self.detectors if x.label == key), None)
        return (not own.enabled) if own else key in self.disabledDetectors

    def policy(self, pid: str) -> PolicySummary | None:
        return next((p for p in self.allPolicies if p.id == pid), None)

    def group(self, gid: str) -> Group | None:
        return next((g for g in self.groups if g.id == gid), None)

    def rule(self, rid: str) -> Rule | None:
        return next((r for r in self.rules if r.id == rid), None)

    def block_list(self, lid: str) -> BlockList | None:
        return next((b for b in self.blockLists if b.id == lid), None)

    def connectors_for(self, channel: str) -> list[Connector]:
        return self.connectors.get(channel, [])

    def web_category(self, name: str) -> WebCategory | None:
        return next((w for w in self.webCategories if w.name.lower() == name.lower()), None)


_LABEL = re.compile(r'^[A-Z][A-Z0-9_]{1,63}$')
_CONTROL = re.compile(r'^(agentControl|extensionControl)\.[A-Za-z][A-Za-z0-9]{1,63}$')


def _s(v: Any, limit: int = 300) -> str:
    return v.strip()[:limit] if isinstance(v, str) else ''


def _strs(v: Any, limit: int = 500) -> list[str]:
    if not isinstance(v, list):
        return []
    out: list[str] = []
    for x in v[:limit]:
        s = _s(x)
        if s and s not in out:
            out.append(s)
    return out


def _b(v: Any) -> bool:
    return v is True


def _dicts(v: Any, limit: int = 2000) -> list[dict]:
    return [x for x in v[:limit] if isinstance(x, dict)] if isinstance(v, list) else []


def load_context(raw: Any) -> PlatformContext:
    """Build a context from untrusted JSON. Never raises."""
    raw = raw if isinstance(raw, dict) else {}
    ctx = PlatformContext()

    for g in _dicts(raw.get('groups')):
        gid, name = _s(g.get('id'), 64), _s(g.get('name'), 120)
        if gid and name and gid != 'default' and not ctx.group(gid):
            ctx.groups.append(Group(gid, name))

    seen: set[str] = set()
    for r in _dicts(raw.get('detectionRules')):
        rid, name = _s(r.get('id'), 64), _s(r.get('name'), 200)
        if not rid or not name or rid in seen:
            continue
        sev = _s(r.get('severity')).upper()
        ctx.rules.append(Rule(
            id=rid, name=name,
            entities=[e.upper() for e in _strs(r.get('entities'))],
            keywords=[k.lower() for k in _strs(r.get('keywords'))],
            fileTypes=[f.lower() for f in _strs(r.get('fileTypes'))],
            severity=sev if sev in ('LOW', 'MEDIUM', 'HIGH') else 'HIGH',
            builtin=_b(r.get('builtin')) or rid.startswith('builtin_'),
        ))
        seen.add(rid)
    # The server seeds the presets on every boot, so they always exist. Add any the caller left out.
    for r in catalog.builtin_rules():
        if r['id'] not in seen:
            ctx.rules.append(Rule(r['id'], r['name'], list(r['entities']), list(r['keywords']),
                                  list(r['fileTypes']), r['severity'], True))

    seen = set()
    for b in _dicts(raw.get('blockLists')):
        lid, name, kind = _s(b.get('id'), 64), _s(b.get('name'), 200), _s(b.get('kind'))
        if not lid or not name or kind not in ('web', 'app') or lid in seen:
            continue
        ctx.blockLists.append(BlockList(lid, name, kind, [e.lower() for e in _strs(b.get('entries'), 5000)],
                                        _b(b.get('builtin')) or lid.startswith('builtin_')))
        seen.add(lid)
    for b in catalog.builtin_block_lists():
        if b['id'] not in seen:
            ctx.blockLists.append(BlockList(b['id'], b['name'], b['kind'], list(b['entries']), True))

    if isinstance(raw.get('webCategories'), list):
        ctx.webCategoriesKnown = True
    for w in _dicts(raw.get('webCategories')):
        name = _s(w.get('name'), 120)
        if name and not ctx.web_category(name):
            n = w.get('domains')
            ctx.webCategories.append(WebCategory(name, n if isinstance(n, int) and n >= 0 else 0))

    conns = raw.get('connectors')
    if isinstance(conns, dict):
        ctx.connectorsKnown = True
        for ch in CONNECTOR_CHANNELS:
            for c in _dicts(conns.get(ch)):
                cid = _s(c.get('id'), 64)
                if not cid:
                    continue
                ctx.connectors.setdefault(ch, []).append(Connector(
                    id=cid, name=_s(c.get('name'), 200) or _s(c.get('teamName'), 200) or cid,
                    hasAdminToken=_b(c.get('hasAdminToken')),
                    remediationEnabled=_b(c.get('remediationEnabled')),
                    provider=_s(c.get('provider'), 16) if ch == 'email' else '',
                ))

    builtin_keys = catalog.entity_keys()
    for x in _dicts(raw.get('detectors'), 500):
        label, name = _s(x.get('label'), 64).upper(), _s(x.get('name'), 120)
        # A custom detector may not shadow a built-in key: the platform refuses that too.
        if _LABEL.match(label) and label not in builtin_keys and not any(d.label == label for d in ctx.detectors):
            ctx.detectors.append(Detector(label, name or label.replace('_', ' ').title(), x.get('enabled') is not False))
    ctx.disabledDetectors = {k.upper() for k in _strs(raw.get('disabledDetectors'))}

    if isinstance(raw.get('policies'), list):
        ctx.policiesKnown = True
    for p in _dicts(raw.get('policies')):
        pid = _s(p.get('id'), 64)
        if _s(p.get('name'), 200):
            ctx.policyNames.add(_s(p.get('name'), 200).lower())
        if not pid or any(x.id == pid for x in ctx.allPolicies):
            continue
        ptype = _s(p.get('type')) or _type_of(p)
        if ptype not in catalog.POLICY_TYPES:
            continue
        summary = PolicySummary(
            id=pid, name=_s(p.get('name'), 200) or pid, type=ptype,
            groupId=_s(p.get('groupId'), 64) or 'default',
            action=_s(p.get('action')) or 'warn',
            connectorIds=_strs(p.get('connectorIds')),
            provider=_email_provider(p) if ptype == 'email' else '',
            active=p.get('active') is not False,
            controls={c for c in _strs(p.get('controls'), 200) if _CONTROL.match(c)},
        )
        if isinstance(p.get('controls'), list):
            ctx.controlsKnown = True
        if isinstance(p.get('blockedDomains'), list):
            ctx.blocksKnown = True
            summary.blockedDomains = {x.lower() for x in _strs(p.get('blockedDomains'), 500)}
            summary.blockedApps = {x.lower() for x in _strs(p.get('blockedApplications'), 500)}
            summary.webLists = _strs(p.get('webBlockListIds'), 100)
            summary.appLists = _strs(p.get('appBlockListIds'), 100)
        ctx.allPolicies.append(summary)
        if summary.active:
            ctx.policies.append(summary)
    return ctx


def _email_provider(p: dict) -> str:
    """As the email form's emailProviderOf(): 'm365' only when saved so, otherwise the Gmail policy."""
    if p.get('provider') == 'm365':
        return 'm365'
    ch = p.get('channels')
    em = ch.get('email') if isinstance(ch, dict) else None
    return 'm365' if isinstance(em, dict) and em.get('provider') == 'm365' else 'google'


def _type_of(p: dict) -> str:
    """A full policy row (as GET /api/admin/policies returns it) has no type field: derive it."""
    ch = p.get('channels')
    unknown = False
    if isinstance(ch, dict):
        for k, v in ch.items():
            if isinstance(v, dict) and v.get('enabled') is True:
                if k in catalog.CHANNEL_TYPES:
                    return k
                unknown = True
    # A channel the copilot does not know yet (a connector added to the platform later):
    # leave the policy out rather than mistake it for an endpoint policy.
    return '' if unknown else 'endpoint'

	"""PlatformContext: what the tenant looks like right now.

	The Node proxy builds this from the database on every request (groups, rules, block lists,
	connectors, active policies) and the copilot never reads the database itself. Loading is
	deliberately forgiving: a malformed entry is dropped, never raised, because a bad context
	must degrade answers, not crash the service.
	"""
	from __future__ import annotations

	import re
	from dataclasses import dataclass, field
	from typing import Any

	from . import catalog

	CONNECTOR_CHANNELS = ('email', 'slack', 'teams', 'drive', 'chat', 'm365files', 's3', 'gcs', 'azure', 'github')


	@dataclass
	class Group:
	id: str
	name: str


	@dataclass
	class Rule:
	id: str
	name: str
	entities: list[str]
	keywords: list[str]
	fileTypes: list[str]
	severity: str
	builtin: bool


	@dataclass
	class BlockList:
	id: str
	name: str
	kind: str # 'web' \| 'app'
	entries: list[str]
	builtin: bool


	@dataclass
	class WebCategory:
	name: str # as the admin named it in Web Filtering; policies store this name
	domains: int # how many sites it holds (0 until it has been filled)


	@dataclass
	class Connector:
	id: str
	name: str
	hasAdminToken: bool = False # slack: needed for the delete action
	remediationEnabled: bool = False # email: needed for the remediate action
	provider: str = '' # email: 'google' \| 'm365'


	@dataclass
	class Detector:
	"""An admin-made regex detector (Policies -> Detectors). Rules reference it by label."""
	label: str # CUSTOMER_ID: what a detection rule stores in `entities`
	name: str # "Customer ID": what the admin calls it
	enabled: bool


	@dataclass
	class PolicySummary:
	id: str
	name: str
	type: str
	groupId: str
	action: str
	connectorIds: list[str]
	provider: str = '' # email policies: 'google' \| 'm365' (one saved before the split is Gmail)
	active: bool = True
	# Device and browser switches that are on ("agentControl.blockPrint"), for unblock requests.
	controls: set[str] = field(default_factory=set)
	# Sites and apps it blocks by name, and the block lists it uses: for "unblock chatgpt".
	blockedDomains: set[str] = field(default_factory=set)
	blockedApps: set[str] = field(default_factory=set)
	webLists: list[str] = field(default_factory=list)
	appLists: list[str] = field(default_factory=list)


	@dataclass
	class PlatformContext:
	groups: list[Group] = field(default_factory=list)
	rules: list[Rule] = field(default_factory=list)
	blockLists: list[BlockList] = field(default_factory=list)
	connectors: dict[str, list[Connector]] = field(default_factory=dict)
	# Active policies only. The copilot uses them to explain how a new one will combine.
	policies: list[PolicySummary] = field(default_factory=list)
	# Every policy, on or off: what a request may add to instead of creating another one.
	allPolicies: list[PolicySummary] = field(default_factory=list)
	# Names of every policy, active or not, so a new one never repeats a name.
	policyNames: set[str] = field(default_factory=set)
	# True when the caller told us which connectors exist. When it did not, we cannot
	# say "no connector", so connector checks are skipped rather than guessed.
	connectorsKnown: bool = False
	policiesKnown: bool = False
	# Web Filtering categories (tenant-defined). Unknown when the caller sent none.
	webCategories: list[WebCategory] = field(default_factory=list)
	webCategoriesKnown: bool = False
	# True when the caller sent each policy's switches (PolicySummary.controls).
	controlsKnown: bool = False
	# True when the caller sent what each policy blocks (sites, apps, lists).
	blocksKnown: bool = False
	# The tenant's own regex detectors, and the built-in ones an admin switched off.
	detectors: list[Detector] = field(default_factory=list)
	disabledDetectors: set[str] = field(default_factory=set)

	def entity_keys(self) -> frozenset[str]:
	"""Every key a rule may hold for this tenant: the platform's plus its custom detectors."""
	if not self.detectors:
	return catalog.entity_keys()
	return catalog.entity_keys() \| {x.label for x in self.detectors}

	def entity_label(self, key: str) -> str:
	own = next((x for x in self.detectors if x.label == key), None)
	return own.name if own else catalog.entity_label(key)

	def detector_off(self, key: str) -> bool:
	"""Switched off on the Detectors page: a rule may list it, but it finds nothing."""
	own = next((x for x in self.detectors if x.label == key), None)
	return (not own.enabled) if own else key in self.disabledDetectors

	def policy(self, pid: str) -> PolicySummary \| None:
	return next((p for p in self.allPolicies if p.id == pid), None)

	def group(self, gid: str) -> Group \| None:
	return next((g for g in self.groups if g.id == gid), None)

	def rule(self, rid: str) -> Rule \| None:
	return next((r for r in self.rules if r.id == rid), None)

	def block_list(self, lid: str) -> BlockList \| None:
	return next((b for b in self.blockLists if b.id == lid), None)

	def connectors_for(self, channel: str) -> list[Connector]:
	return self.connectors.get(channel, [])

	def web_category(self, name: str) -> WebCategory \| None:
	return next((w for w in self.webCategories if w.name.lower() == name.lower()), None)


	_LABEL = re.compile(r'^[A-Z][A-Z0-9_]{1,63}$')
	_CONTROL = re.compile(r'^(agentControl\|extensionControl)\.[A-Za-z][A-Za-z0-9]{1,63}$')


	def _s(v: Any, limit: int = 300) -> str:
	return v.strip()[:limit] if isinstance(v, str) else ''


	def _strs(v: Any, limit: int = 500) -> list[str]:
	if not isinstance(v, list):
	return []
	out: list[str] = []
	for x in v[:limit]:
	s = _s(x)
	if s and s not in out:
	out.append(s)
	return out


	def _b(v: Any) -> bool:
	return v is True


	def _dicts(v: Any, limit: int = 2000) -> list[dict]:
	return [x for x in v[:limit] if isinstance(x, dict)] if isinstance(v, list) else []


	def load_context(raw: Any) -> PlatformContext:
	"""Build a context from untrusted JSON. Never raises."""
	raw = raw if isinstance(raw, dict) else {}
	ctx = PlatformContext()

	for g in _dicts(raw.get('groups')):
	gid, name = _s(g.get('id'), 64), _s(g.get('name'), 120)
	if gid and name and gid != 'default' and not ctx.group(gid):
	ctx.groups.append(Group(gid, name))

	seen: set[str] = set()
	for r in _dicts(raw.get('detectionRules')):
	rid, name = _s(r.get('id'), 64), _s(r.get('name'), 200)
	if not rid or not name or rid in seen:
	continue
	sev = _s(r.get('severity')).upper()
	ctx.rules.append(Rule(
	id=rid, name=name,
	entities=[e.upper() for e in _strs(r.get('entities'))],
	keywords=[k.lower() for k in _strs(r.get('keywords'))],
	fileTypes=[f.lower() for f in _strs(r.get('fileTypes'))],
	severity=sev if sev in ('LOW', 'MEDIUM', 'HIGH') else 'HIGH',
	builtin=_b(r.get('builtin')) or rid.startswith('builtin_'),
	))
	seen.add(rid)
	# The server seeds the presets on every boot, so they always exist. Add any the caller left out.
	for r in catalog.builtin_rules():
	if r['id'] not in seen:
	ctx.rules.append(Rule(r['id'], r['name'], list(r['entities']), list(r['keywords']),
	list(r['fileTypes']), r['severity'], True))

	seen = set()
	for b in _dicts(raw.get('blockLists')):
	lid, name, kind = _s(b.get('id'), 64), _s(b.get('name'), 200), _s(b.get('kind'))
	if not lid or not name or kind not in ('web', 'app') or lid in seen:
	continue
	ctx.blockLists.append(BlockList(lid, name, kind, [e.lower() for e in _strs(b.get('entries'), 5000)],
	_b(b.get('builtin')) or lid.startswith('builtin_')))
	seen.add(lid)
	for b in catalog.builtin_block_lists():
	if b['id'] not in seen:
	ctx.blockLists.append(BlockList(b['id'], b['name'], b['kind'], list(b['entries']), True))

	if isinstance(raw.get('webCategories'), list):
	ctx.webCategoriesKnown = True
	for w in _dicts(raw.get('webCategories')):
	name = _s(w.get('name'), 120)
	if name and not ctx.web_category(name):
	n = w.get('domains')
	ctx.webCategories.append(WebCategory(name, n if isinstance(n, int) and n >= 0 else 0))

	conns = raw.get('connectors')
	if isinstance(conns, dict):
	ctx.connectorsKnown = True
	for ch in CONNECTOR_CHANNELS:
	for c in _dicts(conns.get(ch)):
	cid = _s(c.get('id'), 64)
	if not cid:
	continue
	ctx.connectors.setdefault(ch, []).append(Connector(
	id=cid, name=_s(c.get('name'), 200) or _s(c.get('teamName'), 200) or cid,
	hasAdminToken=_b(c.get('hasAdminToken')),
	remediationEnabled=_b(c.get('remediationEnabled')),
	provider=_s(c.get('provider'), 16) if ch == 'email' else '',
	))

	builtin_keys = catalog.entity_keys()
	for x in _dicts(raw.get('detectors'), 500):
	label, name = _s(x.get('label'), 64).upper(), _s(x.get('name'), 120)
	# A custom detector may not shadow a built-in key: the platform refuses that too.
	if _LABEL.match(label) and label not in builtin_keys and not any(d.label == label for d in ctx.detectors):
	ctx.detectors.append(Detector(label, name or label.replace('_', ' ').title(), x.get('enabled') is not False))
	ctx.disabledDetectors = {k.upper() for k in _strs(raw.get('disabledDetectors'))}

	if isinstance(raw.get('policies'), list):
	ctx.policiesKnown = True
	for p in _dicts(raw.get('policies')):
	pid = _s(p.get('id'), 64)
	if _s(p.get('name'), 200):
	ctx.policyNames.add(_s(p.get('name'), 200).lower())
	if not pid or any(x.id == pid for x in ctx.allPolicies):
	continue
	ptype = _s(p.get('type')) or _type_of(p)
	if ptype not in catalog.POLICY_TYPES:
	continue
	summary = PolicySummary(
	id=pid, name=_s(p.get('name'), 200) or pid, type=ptype,
	groupId=_s(p.get('groupId'), 64) or 'default',
	action=_s(p.get('action')) or 'warn',
	connectorIds=_strs(p.get('connectorIds')),
	provider=_email_provider(p) if ptype == 'email' else '',
	active=p.get('active') is not False,
	controls={c for c in _strs(p.get('controls'), 200) if _CONTROL.match(c)},
	)
	if isinstance(p.get('controls'), list):
	ctx.controlsKnown = True
	if isinstance(p.get('blockedDomains'), list):
	ctx.blocksKnown = True
	summary.blockedDomains = {x.lower() for x in _strs(p.get('blockedDomains'), 500)}
	summary.blockedApps = {x.lower() for x in _strs(p.get('blockedApplications'), 500)}
	summary.webLists = _strs(p.get('webBlockListIds'), 100)
	summary.appLists = _strs(p.get('appBlockListIds'), 100)
	ctx.allPolicies.append(summary)
	if summary.active:
	ctx.policies.append(summary)
	return ctx


	def _email_provider(p: dict) -> str:
	"""As the email form's emailProviderOf(): 'm365' only when saved so, otherwise the Gmail policy."""
	if p.get('provider') == 'm365':
	return 'm365'
	ch = p.get('channels')
	em = ch.get('email') if isinstance(ch, dict) else None
	return 'm365' if isinstance(em, dict) and em.get('provider') == 'm365' else 'google'


	def _type_of(p: dict) -> str:
	"""A full policy row (as GET /api/admin/policies returns it) has no type field: derive it."""
	ch = p.get('channels')
	unknown = False
	if isinstance(ch, dict):
	for k, v in ch.items():
	if isinstance(v, dict) and v.get('enabled') is True:
	if k in catalog.CHANNEL_TYPES:
	return k
	unknown = True
	# A channel the copilot does not know yet (a connector added to the platform later):
	# leave the policy out rather than mistake it for an endpoint policy.
	return '' if unknown else 'endpoint'