Blame for copilot/platform/context.py

"""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'