rule_engine.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313
  1. #!/usr/bin/env python3
  2. """Rule evaluation engine for hookify plugin."""
  3. import re
  4. import sys
  5. from functools import lru_cache
  6. from typing import List, Dict, Any, Optional
  7. # Import from local module
  8. from core.config_loader import Rule, Condition
  9. # Cache compiled regexes (max 128 patterns)
  10. @lru_cache(maxsize=128)
  11. def compile_regex(pattern: str) -> re.Pattern:
  12. """Compile regex pattern with caching.
  13. Args:
  14. pattern: Regex pattern string
  15. Returns:
  16. Compiled regex pattern
  17. """
  18. return re.compile(pattern, re.IGNORECASE)
  19. class RuleEngine:
  20. """Evaluates rules against hook input data."""
  21. def __init__(self):
  22. """Initialize rule engine."""
  23. # No need for instance cache anymore - using global lru_cache
  24. pass
  25. def evaluate_rules(self, rules: List[Rule], input_data: Dict[str, Any]) -> Dict[str, Any]:
  26. """Evaluate all rules and return combined results.
  27. Checks all rules and accumulates matches. Blocking rules take priority
  28. over warning rules. All matching rule messages are combined.
  29. Args:
  30. rules: List of Rule objects to evaluate
  31. input_data: Hook input JSON (tool_name, tool_input, etc.)
  32. Returns:
  33. Response dict with systemMessage, hookSpecificOutput, etc.
  34. Empty dict {} if no rules match.
  35. """
  36. hook_event = input_data.get('hook_event_name', '')
  37. blocking_rules = []
  38. warning_rules = []
  39. for rule in rules:
  40. if self._rule_matches(rule, input_data):
  41. if rule.action == 'block':
  42. blocking_rules.append(rule)
  43. else:
  44. warning_rules.append(rule)
  45. # If any blocking rules matched, block the operation
  46. if blocking_rules:
  47. messages = [f"**[{r.name}]**\n{r.message}" for r in blocking_rules]
  48. combined_message = "\n\n".join(messages)
  49. # Use appropriate blocking format based on event type
  50. if hook_event == 'Stop':
  51. return {
  52. "decision": "block",
  53. "reason": combined_message,
  54. "systemMessage": combined_message
  55. }
  56. elif hook_event in ['PreToolUse', 'PostToolUse']:
  57. return {
  58. "hookSpecificOutput": {
  59. "hookEventName": hook_event,
  60. "permissionDecision": "deny"
  61. },
  62. "systemMessage": combined_message
  63. }
  64. else:
  65. # For other events, just show message
  66. return {
  67. "systemMessage": combined_message
  68. }
  69. # If only warnings, show them but allow operation
  70. if warning_rules:
  71. messages = [f"**[{r.name}]**\n{r.message}" for r in warning_rules]
  72. return {
  73. "systemMessage": "\n\n".join(messages)
  74. }
  75. # No matches - allow operation
  76. return {}
  77. def _rule_matches(self, rule: Rule, input_data: Dict[str, Any]) -> bool:
  78. """Check if rule matches input data.
  79. Args:
  80. rule: Rule to evaluate
  81. input_data: Hook input data
  82. Returns:
  83. True if rule matches, False otherwise
  84. """
  85. # Extract tool information
  86. tool_name = input_data.get('tool_name', '')
  87. tool_input = input_data.get('tool_input', {})
  88. # Check tool matcher if specified
  89. if rule.tool_matcher:
  90. if not self._matches_tool(rule.tool_matcher, tool_name):
  91. return False
  92. # If no conditions, don't match
  93. # (Rules must have at least one condition to be valid)
  94. if not rule.conditions:
  95. return False
  96. # All conditions must match
  97. for condition in rule.conditions:
  98. if not self._check_condition(condition, tool_name, tool_input, input_data):
  99. return False
  100. return True
  101. def _matches_tool(self, matcher: str, tool_name: str) -> bool:
  102. """Check if tool_name matches the matcher pattern.
  103. Args:
  104. matcher: Pattern like "Bash", "Edit|Write", "*"
  105. tool_name: Actual tool name
  106. Returns:
  107. True if matches
  108. """
  109. if matcher == '*':
  110. return True
  111. # Split on | for OR matching
  112. patterns = matcher.split('|')
  113. return tool_name in patterns
  114. def _check_condition(self, condition: Condition, tool_name: str,
  115. tool_input: Dict[str, Any], input_data: Dict[str, Any] = None) -> bool:
  116. """Check if a single condition matches.
  117. Args:
  118. condition: Condition to check
  119. tool_name: Tool being used
  120. tool_input: Tool input dict
  121. input_data: Full hook input data (for Stop events, etc.)
  122. Returns:
  123. True if condition matches
  124. """
  125. # Extract the field value to check
  126. field_value = self._extract_field(condition.field, tool_name, tool_input, input_data)
  127. if field_value is None:
  128. return False
  129. # Apply operator
  130. operator = condition.operator
  131. pattern = condition.pattern
  132. if operator == 'regex_match':
  133. return self._regex_match(pattern, field_value)
  134. elif operator == 'contains':
  135. return pattern in field_value
  136. elif operator == 'equals':
  137. return pattern == field_value
  138. elif operator == 'not_contains':
  139. return pattern not in field_value
  140. elif operator == 'starts_with':
  141. return field_value.startswith(pattern)
  142. elif operator == 'ends_with':
  143. return field_value.endswith(pattern)
  144. else:
  145. # Unknown operator
  146. return False
  147. def _extract_field(self, field: str, tool_name: str,
  148. tool_input: Dict[str, Any], input_data: Dict[str, Any] = None) -> Optional[str]:
  149. """Extract field value from tool input or hook input data.
  150. Args:
  151. field: Field name like "command", "new_text", "file_path", "reason", "transcript"
  152. tool_name: Tool being used (may be empty for Stop events)
  153. tool_input: Tool input dict
  154. input_data: Full hook input (for accessing transcript_path, reason, etc.)
  155. Returns:
  156. Field value as string, or None if not found
  157. """
  158. # Direct tool_input fields
  159. if field in tool_input:
  160. value = tool_input[field]
  161. if isinstance(value, str):
  162. return value
  163. return str(value)
  164. # For Stop events and other non-tool events, check input_data
  165. if input_data:
  166. # Stop event specific fields
  167. if field == 'reason':
  168. return input_data.get('reason', '')
  169. elif field == 'transcript':
  170. # Read transcript file if path provided
  171. transcript_path = input_data.get('transcript_path')
  172. if transcript_path:
  173. try:
  174. with open(transcript_path, 'r') as f:
  175. return f.read()
  176. except FileNotFoundError:
  177. print(f"Warning: Transcript file not found: {transcript_path}", file=sys.stderr)
  178. return ''
  179. except PermissionError:
  180. print(f"Warning: Permission denied reading transcript: {transcript_path}", file=sys.stderr)
  181. return ''
  182. except (IOError, OSError) as e:
  183. print(f"Warning: Error reading transcript {transcript_path}: {e}", file=sys.stderr)
  184. return ''
  185. except UnicodeDecodeError as e:
  186. print(f"Warning: Encoding error in transcript {transcript_path}: {e}", file=sys.stderr)
  187. return ''
  188. elif field == 'user_prompt':
  189. # For UserPromptSubmit events
  190. return input_data.get('user_prompt', '')
  191. # Handle special cases by tool type
  192. if tool_name == 'Bash':
  193. if field == 'command':
  194. return tool_input.get('command', '')
  195. elif tool_name in ['Write', 'Edit']:
  196. if field == 'content':
  197. # Write uses 'content', Edit has 'new_string'
  198. return tool_input.get('content') or tool_input.get('new_string', '')
  199. elif field == 'new_text' or field == 'new_string':
  200. return tool_input.get('new_string', '')
  201. elif field == 'old_text' or field == 'old_string':
  202. return tool_input.get('old_string', '')
  203. elif field == 'file_path':
  204. return tool_input.get('file_path', '')
  205. elif tool_name == 'MultiEdit':
  206. if field == 'file_path':
  207. return tool_input.get('file_path', '')
  208. elif field in ['new_text', 'content']:
  209. # Concatenate all edits
  210. edits = tool_input.get('edits', [])
  211. return ' '.join(e.get('new_string', '') for e in edits)
  212. return None
  213. def _regex_match(self, pattern: str, text: str) -> bool:
  214. """Check if pattern matches text using regex.
  215. Args:
  216. pattern: Regex pattern
  217. text: Text to match against
  218. Returns:
  219. True if pattern matches
  220. """
  221. try:
  222. # Use cached compiled regex (LRU cache with max 128 patterns)
  223. regex = compile_regex(pattern)
  224. return bool(regex.search(text))
  225. except re.error as e:
  226. print(f"Invalid regex pattern '{pattern}': {e}", file=sys.stderr)
  227. return False
  228. # For testing
  229. if __name__ == '__main__':
  230. from core.config_loader import Condition, Rule
  231. # Test rule evaluation
  232. rule = Rule(
  233. name="test-rm",
  234. enabled=True,
  235. event="bash",
  236. conditions=[
  237. Condition(field="command", operator="regex_match", pattern=r"rm\s+-rf")
  238. ],
  239. message="Dangerous rm command!"
  240. )
  241. engine = RuleEngine()
  242. # Test matching input
  243. test_input = {
  244. "tool_name": "Bash",
  245. "tool_input": {
  246. "command": "rm -rf /tmp/test"
  247. }
  248. }
  249. result = engine.evaluate_rules([rule], test_input)
  250. print("Match result:", result)
  251. # Test non-matching input
  252. test_input2 = {
  253. "tool_name": "Bash",
  254. "tool_input": {
  255. "command": "ls -la"
  256. }
  257. }
  258. result2 = engine.evaluate_rules([rule], test_input2)
  259. print("Non-match result:", result2)