main
py 331 lines 12.7 KB
Raw
1 """
2 Service layer for the CoPilot Case-Template **Library** — a read-only catalog
3 of investigation playbooks authored as YAML in
4 https://github.com/socfortress/CoPilot-Case-Templates.
5
6 The Library is *not* a second template-management system. It is a fetcher +
7 parser + cache that surfaces YAML-defined playbooks to the admin UI so they
8 can be imported into the existing ``CaseTemplate`` tables via
9 ``create_template``. Once imported, an entry becomes a normal DB row and is
10 managed through the existing CRUD endpoints — edits in the UI never flow
11 back to GitHub, and changes pushed to GitHub never retroactively update
12 already-imported templates.
13
14 Conventions mirrored from
15 ``app.integrations.copilot_searches.services.copilot_searches.RulesCache``:
16 - In-process cache with a TTL (30 minutes by default).
17 - ``asyncio.Lock`` around the refresh so concurrent requests don't fire
18 duplicate fetches.
19 - Best-effort YAML validation per file; one bad file does not fail the
20 whole library.
21 - ``GITHUB_TOKEN`` env var (if present) is sent as a Bearer header to dodge
22 GitHub's 60/hr unauthenticated rate limit. Same env var the CoPilot
23 Searches loader uses.
24
25 Library YAML schema (see ``SCHEMA.md`` in the playbook repo):
26
27 key: str (required, unique across the repo)
28 name: str (required, <= 255 chars)
29 description: str (optional)
30 source: str (optional, <= 50 chars)
31 match: (optional; both keys required when the block is present)
32 field: str (e.g., "data_win_system_eventID")
33 value: str (equality target; stored as-is)
34 tags: dict (optional, library-only metadata, not persisted)
35 tasks:
36 - title: str (required, <= 500 chars)
37 description: str (optional)
38 guidelines: str (optional)
39 mandatory: bool (optional, default False)
40 order_index: int (required, >= 0)
41 """
42
43 import asyncio
44 from datetime import datetime
45 from datetime import timedelta
46 from typing import Any
47 from typing import Dict
48 from typing import List
49 from typing import Optional
50
51 import httpx
52 import yaml
53 from loguru import logger
54
55 # Where the playbook YAMLs live. Owned by SOCFortress; PR-able by anyone with
56 # repo access. CoPilot only ever reads from this repo — never writes.
57 GITHUB_REPO = "socfortress/CoPilot-Case-Templates"
58 GITHUB_BRANCH = "main"
59 GITHUB_API_BASE = "https://api.github.com"
60 GITHUB_RAW_BASE = "https://raw.githubusercontent.com"
61
62 # Only files matching this prefix-set are treated as library entries. Future
63 # domain folders (e.g. "linux/", "office365/") can be added here without code
64 # changes elsewhere.
65 LIBRARY_PATH_PREFIXES = ("sysmon/",)
66
67 # In-memory cache TTL. Matches the CoPilot Searches cache so the operator's
68 # mental model for "how stale can this be" is the same.
69 CACHE_TTL_MINUTES = 30
70
71
72 def _github_headers() -> Dict[str, str]:
73 """
74 Headers sent on every GitHub call. The ``Authorization`` header is added
75 only when ``GITHUB_TOKEN`` is set — its absence falls back to the 60/hr
76 unauthenticated quota, which is enough for casual use but trips quickly
77 on shared dev environments. Same convention as the CoPilot Searches loader.
78 """
79 headers = {"Accept": "application/vnd.github+json"}
80 return headers
81
82
83 class TemplateLibraryCache:
84 """In-memory cache of parsed YAML library entries.
85
86 The cache populates lazily on first access (``ensure_loaded``) and refreshes
87 after the TTL expires or when ``refresh`` is invoked explicitly (e.g. via
88 the ``POST /library/refresh`` admin endpoint).
89 """
90
91 def __init__(self) -> None:
92 self._entries: Dict[str, Dict[str, Any]] = {} # key -> parsed dict
93 self._invalid_paths: List[str] = [] # paths skipped during last refresh
94 self._last_refresh: Optional[datetime] = None
95 self._lock = asyncio.Lock()
96
97 # ----- lifecycle -----
98
99 @property
100 def is_stale(self) -> bool:
101 if self._last_refresh is None:
102 return True
103 return datetime.utcnow() - self._last_refresh > timedelta(minutes=CACHE_TTL_MINUTES)
104
105 async def ensure_loaded(self) -> None:
106 if self.is_stale:
107 await self.refresh()
108
109 async def refresh(self) -> int:
110 """
111 Re-fetch the repo tree and every YAML under a recognised prefix.
112
113 Returns the number of valid entries loaded. Bad YAML files are logged
114 and skipped (their paths are recorded in ``invalid_paths``); the
115 whole-library refresh never raises on per-file parse errors so a
116 single malformed file in the upstream repo doesn't take the feature
117 down. Network failures DO bubble up — those should reach the operator.
118 """
119 async with self._lock:
120 logger.info(
121 f"Refreshing case-template library from github.com/{GITHUB_REPO}@{GITHUB_BRANCH}",
122 )
123
124 entries: Dict[str, Dict[str, Any]] = {}
125 invalid: List[str] = []
126
127 async with httpx.AsyncClient(timeout=30.0, headers=_github_headers()) as client:
128 tree_url = f"{GITHUB_API_BASE}/repos/{GITHUB_REPO}/git/trees/{GITHUB_BRANCH}?recursive=1"
129 response = await client.get(tree_url)
130 response.raise_for_status()
131 tree = response.json()
132
133 yaml_paths = [
134 item["path"]
135 for item in tree.get("tree", [])
136 if item.get("type") == "blob"
137 and item.get("path", "").endswith((".yaml", ".yml"))
138 and any(item["path"].startswith(p) for p in LIBRARY_PATH_PREFIXES)
139 ]
140 logger.info(f"Library tree has {len(yaml_paths)} YAML file(s) to parse")
141
142 # Fan out the per-file fetches; gather lets one slow file
143 # block the others without serialising the whole pull.
144 tasks = [self._fetch_and_parse(client, p) for p in yaml_paths]
145 results = await asyncio.gather(*tasks, return_exceptions=True)
146
147 for path, result in zip(yaml_paths, results):
148 if isinstance(result, Exception):
149 logger.warning(f"Library: failed to fetch/parse {path}: {result}")
150 invalid.append(path)
151 continue
152 if result is None:
153 invalid.append(path)
154 continue
155
156 entry = result
157 key = entry.get("key")
158 if not key:
159 logger.warning(f"Library: {path} missing 'key' field, skipping")
160 invalid.append(path)
161 continue
162 if key in entries:
163 logger.warning(
164 f"Library: duplicate key '{key}' in {path} " f"(already used by {entries[key].get('_file_path')}), skipping",
165 )
166 invalid.append(path)
167 continue
168 entry["_file_path"] = path
169 entries[key] = entry
170
171 self._entries = entries
172 self._invalid_paths = invalid
173 self._last_refresh = datetime.utcnow()
174 logger.info(
175 f"Library loaded: {len(entries)} valid entr(ies), {len(invalid)} skipped",
176 )
177 return len(entries)
178
179 async def _fetch_and_parse(
180 self,
181 client: httpx.AsyncClient,
182 path: str,
183 ) -> Optional[Dict[str, Any]]:
184 """Fetch one YAML file, parse it, and run light shape validation."""
185 raw_url = f"{GITHUB_RAW_BASE}/{GITHUB_REPO}/{GITHUB_BRANCH}/{path}"
186 response = await client.get(raw_url)
187 response.raise_for_status()
188 try:
189 data = yaml.safe_load(response.text)
190 except yaml.YAMLError as e:
191 logger.warning(f"Library: YAML parse error in {path}: {e}")
192 return None
193
194 if not isinstance(data, dict):
195 logger.warning(f"Library: {path} top-level is not a mapping, skipping")
196 return None
197
198 return _normalize_entry(data)
199
200 # ----- read accessors -----
201
202 @property
203 def entries(self) -> Dict[str, Dict[str, Any]]:
204 return self._entries
205
206 @property
207 def invalid_paths(self) -> List[str]:
208 return list(self._invalid_paths)
209
210 @property
211 def last_refresh(self) -> Optional[datetime]:
212 return self._last_refresh
213
214 def get_entry(self, key: str) -> Optional[Dict[str, Any]]:
215 return self._entries.get(key)
216
217
218 # ---------------------------------------------------------------------------
219 # Normalisation + validation helpers
220 # ---------------------------------------------------------------------------
221
222
223 def _normalize_entry(data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
224 """
225 Coerce a raw YAML dict into the canonical library-entry shape. Returns
226 ``None`` if the entry is missing required fields, so the caller can skip
227 it cleanly. Optional fields are normalised to consistent defaults so the
228 downstream Pydantic response model never sees ``None``-vs-missing
229 ambiguity.
230
231 Required: ``key``, ``name``, ``tasks`` (may be empty).
232 """
233 key = data.get("key")
234 name = data.get("name")
235 if not isinstance(key, str) or not key.strip():
236 return None
237 if not isinstance(name, str) or not name.strip():
238 return None
239
240 raw_tasks = data.get("tasks") or []
241 if not isinstance(raw_tasks, list):
242 return None
243
244 normalised_tasks: List[Dict[str, Any]] = []
245 for raw in raw_tasks:
246 if not isinstance(raw, dict):
247 continue
248 title = raw.get("title")
249 order_index = raw.get("order_index")
250 if not isinstance(title, str) or not title.strip():
251 continue
252 if not isinstance(order_index, int) or order_index < 0:
253 continue
254 normalised_tasks.append(
255 {
256 "title": title.strip(),
257 "description": raw.get("description") or None,
258 "guidelines": raw.get("guidelines") or None,
259 "mandatory": bool(raw.get("mandatory", False)),
260 "order_index": order_index,
261 },
262 )
263
264 # Sort tasks by their declared order_index up front so import order is
265 # deterministic regardless of how the YAML happened to list them.
266 normalised_tasks.sort(key=lambda t: t["order_index"])
267
268 tags = data.get("tags")
269 if not isinstance(tags, dict):
270 tags = {}
271
272 # Optional ``match:`` block for conditional auto-apply. Both ``field`` and
273 # ``value`` must be present (and non-empty strings) or the block is
274 # ignored — half-set conditions never trigger, so we'd rather drop the
275 # block and log than persist a misleading template.
276 match_field: Optional[str] = None
277 match_value: Optional[str] = None
278 match_block = data.get("match")
279 if isinstance(match_block, dict):
280 raw_field = match_block.get("field")
281 raw_value = match_block.get("value")
282 if isinstance(raw_field, str) and raw_field.strip() and raw_value is not None:
283 match_field = raw_field.strip()
284 # Coerce to string so YAML-typed values (ints, bools) survive the
285 # round-trip the same way they arrive from OpenSearch — Wazuh
286 # field values come back as strings ("1" not 1).
287 match_value = str(raw_value)
288
289 return {
290 "key": key.strip(),
291 "name": name.strip(),
292 "description": data.get("description") or None,
293 "source": data.get("source") or None,
294 "match_field": match_field,
295 "match_value": match_value,
296 "tags": tags,
297 "tasks": normalised_tasks,
298 }
299
300
301 # ---------------------------------------------------------------------------
302 # Module-level singleton + service functions consumed by the routes.
303 # ---------------------------------------------------------------------------
304
305 template_library_cache = TemplateLibraryCache()
306
307
308 async def list_library_entries() -> List[Dict[str, Any]]:
309 """Return all library entries currently in cache. Loads on first call."""
310 await template_library_cache.ensure_loaded()
311 # Sort for a stable display order in the UI; primary key on `key`.
312 return sorted(template_library_cache.entries.values(), key=lambda e: e["key"])
313
314
315 async def get_library_entry(key: str) -> Optional[Dict[str, Any]]:
316 """Fetch one library entry by its YAML ``key``. Loads on first call."""
317 await template_library_cache.ensure_loaded()
318 return template_library_cache.get_entry(key)
319
320
321 async def refresh_library() -> Dict[str, Any]:
322 """
323 Force a re-fetch of the library repo. Returns a small status payload that
324 the route layer wraps as the HTTP response.
325 """
326 count = await template_library_cache.refresh()
327 return {
328 "loaded": count,
329 "invalid_paths": template_library_cache.invalid_paths,
330 "last_refresh": template_library_cache.last_refresh,
331 }