44Provides comprehensive audit trail for security and compliance.
55"""
66
7+ import hashlib
78import json
89from datetime import datetime
910from pathlib import Path
10- from typing import Dict , Any , Optional
11- import hashlib
12-
11+ from typing import Any , Dict , Optional
1312
1413class AuditLogger :
1514 """
1615 Log all user actions for audit trail.
17-
16+
1817 Features:
1918 - Immutable event logging
2019 - Cryptographic hashing for integrity
2120 - Support for multiple log formats
2221 - Event categorization and tagging
2322 """
24-
23+
2524 def __init__ (self , log_file : str = "audit.log" ):
2625 self .log_file = Path (log_file )
2726 self .log_file .parent .mkdir (parents = True , exist_ok = True )
28-
27+
2928 # Event types for categorization
3029 self .event_types = {
3130 "authentication" : ["login" , "logout" , "token_refresh" ],
3231 "session_management" : ["create" , "update" , "terminate" ],
3332 "worker_management" : ["add" , "remove" , "sync_model" ],
3433 "configuration" : ["read" , "write" , "backup" , "restore" ],
3534 "inference" : ["start" , "stop" , "benchmark" ],
36- "system" : ["health_check" , "error" , "recovery" ]
35+ "system" : ["health_check" , "error" , "recovery" ],
3736 }
38-
37+
3938 def log_action (
40- self ,
41- action_type : str ,
42- resource : str ,
39+ self ,
40+ action_type : str ,
41+ resource : str ,
4342 details : Dict [str , Any ] = None ,
4443 user_id : str = None ,
45- ip_address : str = None
44+ ip_address : str = None ,
4645 ):
4746 """
4847 Record auditable action.
49-
48+
5049 Args:
5150 action_type: Type of action (e.g., "create", "read")
5251 resource: Resource affected (e.g., "session_abc123")
@@ -59,33 +58,33 @@ def log_action(
5958 resource = resource ,
6059 details = details or {},
6160 user_id = user_id ,
62- ip_address = ip_address
61+ ip_address = ip_address ,
6362 )
64-
63+
6564 # Write to log file (append mode)
66- with open (self .log_file , 'a' ) as f :
67- f .write (json .dumps (event ) + ' \n ' )
68-
65+ with open (self .log_file , "a" ) as f :
66+ f .write (json .dumps (event ) + " \n " )
67+
6968 return event
70-
69+
7170 def _create_event (
7271 self ,
7372 action_type : str ,
7473 resource : str ,
7574 details : Dict [str , Any ],
7675 user_id : str = None ,
77- ip_address : str = None
76+ ip_address : str = None ,
7877 ) -> Dict [str , Any ]:
7978 """Create audit event with all required fields."""
80-
79+
8180 # Categorize event type
8281 category = self ._categorize_event (action_type )
83-
82+
8483 # Create unique event ID
8584 event_id = hashlib .sha256 (
8685 f"{ datetime .now ().isoformat ()} { resource } " .encode ()
8786 ).hexdigest ()[:16 ]
88-
87+
8988 event = {
9089 "event_id" : event_id ,
9190 "timestamp" : datetime .now ().isoformat (),
@@ -94,18 +93,18 @@ def _create_event(
9493 "resource" : resource ,
9594 "details" : details ,
9695 "user_id" : user_id or "anonymous" ,
97- "ip_address" : ip_address or "unknown"
96+ "ip_address" : ip_address or "unknown" ,
9897 }
99-
98+
10099 return event
101-
100+
102101 def _categorize_event (self , action_type : str ) -> str :
103102 """Categorize event based on action type."""
104103 for category , types in self .event_types .items ():
105104 if action_type in types :
106105 return category
107106 return "unknown"
108-
107+
109108 def log_error (self , error : Exception , context : Dict [str , Any ] = None ):
110109 """Log error with full stack trace."""
111110 event = {
@@ -119,130 +118,121 @@ def log_error(self, error: Exception, context: Dict[str, Any] = None):
119118 "details" : {
120119 "error_type" : type (error ).__name__ ,
121120 "error_message" : str (error ),
122- "stack_trace" : self ._get_stack_trace ()
121+ "stack_trace" : self ._get_stack_trace (),
123122 },
124123 "user_id" : None ,
125- "ip_address" : None
124+ "ip_address" : None ,
126125 }
127-
128- with open (self .log_file , 'a' ) as f :
129- f .write (json .dumps (event ) + ' \n ' )
130-
126+
127+ with open (self .log_file , "a" ) as f :
128+ f .write (json .dumps (event ) + " \n " )
129+
131130 def _get_stack_trace (self ) -> str :
132131 """Get current stack trace."""
133132 import traceback
133+
134134 return traceback .format_exc ()
135-
135+
136136 def get_events (
137- self ,
138- event_type : str = None ,
139- resource : str = None ,
140- since : str = None
137+ self , event_type : str = None , resource : str = None , since : str = None
141138 ) -> list :
142139 """
143140 Query audit log for events.
144-
141+
145142 Args:
146143 event_type: Filter by action type
147144 resource: Filter by resource
148145 since: Filter by timestamp (ISO format)
149-
146+
150147 Returns:
151148 List of matching events
152149 """
153150 events = []
154-
155- with open (self .log_file , 'r' ) as f :
151+
152+ with open (self .log_file , "r" ) as f :
156153 for line in f :
157154 try :
158155 event = json .loads (line .strip ())
159-
156+
160157 # Apply filters
161158 if event_type and event .get ("action_type" ) != event_type :
162159 continue
163160 if resource and event .get ("resource" ) != resource :
164161 continue
165162 if since and event .get ("timestamp" , "" ) < since :
166163 continue
167-
164+
168165 events .append (event )
169166 except json .JSONDecodeError :
170167 continue
171-
168+
172169 return events
173-
170+
174171 def get_summary (self ) -> Dict [str , Any ]:
175172 """Get audit log summary statistics."""
176- stats = {
177- "total_events" : 0 ,
178- "by_category" : {},
179- "by_action_type" : {}
180- }
181-
173+ stats = {"total_events" : 0 , "by_category" : {}, "by_action_type" : {}}
174+
182175 try :
183- with open (self .log_file , 'r' ) as f :
176+ with open (self .log_file , "r" ) as f :
184177 for line in f :
185178 try :
186179 event = json .loads (line .strip ())
187180 stats ["total_events" ] += 1
188-
181+
189182 category = event .get ("category" , "unknown" )
190183 action_type = event .get ("action_type" , "unknown" )
191-
192- stats ["by_category" ][category ] = \
184+
185+ stats ["by_category" ][category ] = (
193186 stats ["by_category" ].get (category , 0 ) + 1
194- stats ["by_action_type" ][action_type ] = \
187+ )
188+ stats ["by_action_type" ][action_type ] = (
195189 stats ["by_action_type" ].get (action_type , 0 ) + 1
196-
190+ )
191+
197192 except json .JSONDecodeError :
198193 continue
199194 except FileNotFoundError :
200195 pass
201-
202- return stats
203196
197+ return stats
204198
205199class AuditEventStore :
206200 """In-memory audit event store for real-time queries."""
207-
201+
208202 def __init__ (self , max_events : int = 10000 ):
209203 self .events : list = []
210204 self .max_events = max_events
211-
205+
212206 def add_event (self , event : Dict [str , Any ]):
213207 """Add event to store."""
214208 self .events .append (event )
215-
209+
216210 # Keep only most recent events
217211 if len (self .events ) > self .max_events :
218- self .events = self .events [- self .max_events :]
219-
212+ self .events = self .events [- self .max_events :]
213+
220214 def get_recent_events (self , count : int = 100 ) -> list :
221215 """Get most recent events."""
222216 return self .events [- count :] if len (self .events ) > count else self .events .copy ()
223-
217+
224218 def filter_by_type (self , event_type : str ) -> list :
225219 """Filter events by type."""
226220 return [e for e in self .events if e .get ("action_type" ) == event_type ]
227221
228-
229222if __name__ == "__main__" :
230223 # Test audit logging
231224 logger = AuditLogger ()
232-
225+
233226 # Log some test events
234227 logger .log_action (
235228 action_type = "create" ,
236229 resource = "session_abc123" ,
237230 details = {"model" : "model.onnx" , "devices" : ["gpu_0" ]},
238- user_id = "user1"
231+ user_id = "user1" ,
239232 )
240-
233+
241234 logger .log_action (
242- action_type = "read" ,
243- resource = "config.toml" ,
244- details = {},
245- user_id = "admin"
235+ action_type = "read" , resource = "config.toml" , details = {}, user_id = "admin"
246236 )
247-
237+
248238 print ("Audit log created successfully!" )
0 commit comments