computer-use-ootb-internal 0.0.178__py3-none-any.whl → 0.0.180__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,951 +1,951 @@
1
- # src/computer_use_ootb_internal/guard_service.py
2
- import sys
3
- import os
4
- import time
5
- import logging
6
- import subprocess
7
- import pathlib
8
- import ctypes
9
- import threading # For running server thread
10
- import queue # For queuing commands
11
- import requests # Keep for status reporting back
12
- import servicemanager # From pywin32
13
- import win32serviceutil # From pywin32
14
- import win32service # From pywin32
15
- import win32event # From pywin32
16
- import win32api # From pywin32
17
- import win32process # From pywin32
18
- import win32security # From pywin32
19
- import win32profile # From pywin32
20
- import win32ts # From pywin32 (Terminal Services API)
21
- import win32con
22
- import psutil # For process/user info
23
- from flask import Flask, request, jsonify # For embedded server
24
- from waitress import serve # For serving Flask app
25
- import json # Needed for status reporting
26
-
27
- # --- Configuration ---
28
- _SERVICE_NAME = "OOTBGuardService"
29
- _SERVICE_DISPLAY_NAME = "OOTB Guard Service"
30
- _SERVICE_DESCRIPTION = "Background service for OOTB monitoring and remote management (Server POST mode)."
31
- _PACKAGE_NAME = "computer-use-ootb-internal"
32
- _OOTB_MODULE = "computer_use_ootb_internal.app_teachmode"
33
- # --- Server POST Configuration ---
34
- _LISTEN_HOST = "0.0.0.0" # Listen on all interfaces
35
- _LISTEN_PORT = 14000 # Port for server to POST commands TO
36
- # _SHARED_SECRET = "YOUR_SECRET_HERE" # !! REMOVED !! - No secret check implemented now
37
- # --- End Server POST Configuration ---
38
- _SERVER_STATUS_REPORT_URL = "http://52.160.105.102:7000/guard/status" # URL to POST status back TO (Path changed)
39
- _LOG_FILE = pathlib.Path(os.environ['PROGRAMDATA']) / "OOTBGuardService" / "guard_post_mode.log" # Different log file
40
- # --- End Configuration ---
41
-
42
- _LOG_FILE.parent.mkdir(parents=True, exist_ok=True)
43
- logging.basicConfig(
44
- filename=_LOG_FILE,
45
- level=logging.INFO,
46
- format='%(asctime)s %(levelname)s:%(name)s:%(threadName)s: %(message)s'
47
- )
48
-
49
- # --- Global service instance reference (needed for Flask routes) ---
50
- _service_instance = None
51
-
52
- # --- Flask App Definition ---
53
- flask_app = Flask(__name__)
54
-
55
- @flask_app.route('/command', methods=['POST'])
56
- def receive_command():
57
- global _service_instance
58
- if not _service_instance:
59
- logging.error("Received command but service instance is not set.")
60
- return jsonify({"error": "Service not ready"}), 503
61
-
62
- # --- Authentication REMOVED ---
63
- # secret = request.headers.get('X-Guard-Secret')
64
- # if not secret or secret != _SHARED_SECRET:
65
- # logging.warning(f"Unauthorized command POST received (Invalid/Missing X-Guard-Secret). Remote Addr: {request.remote_addr}")
66
- # return jsonify({"error": "Unauthorized"}), 403
67
- # --- End Authentication REMOVED ---
68
-
69
- if not request.is_json:
70
- logging.warning("Received non-JSON command POST.")
71
- return jsonify({"error": "Request must be JSON"}), 400
72
-
73
- command = request.get_json()
74
- logging.info(f"Received command via POST: {command}")
75
-
76
- # Basic validation
77
- action = command.get("action")
78
- command_id = command.get("command_id", "N/A") # Use for status reporting
79
- if not action:
80
- logging.error(f"Received command POST with no action: {command}")
81
- return jsonify({"error": "Missing 'action' in command"}), 400
82
-
83
- # Queue the command for processing in a background thread
84
- _service_instance.command_queue.put((command_id, command))
85
- logging.info(f"Queued command {command_id} ({action}) for processing.")
86
-
87
- return jsonify({"message": f"Command {command_id} received and queued"}), 202 # Accepted
88
-
89
- @flask_app.route('/internal/user_connected', methods=['POST'])
90
- def user_connected_notification():
91
- """Internal endpoint triggered by the scheduled task helper script."""
92
- global _service_instance
93
- if not _service_instance:
94
- logging.error("Received user_connected signal but service instance is not set.")
95
- return jsonify({"error": "Service not ready"}), 503
96
-
97
- if not request.is_json:
98
- logging.warning("Received non-JSON user_connected signal.")
99
- return jsonify({"error": "Request must be JSON"}), 400
100
-
101
- data = request.get_json()
102
- username = data.get("username")
103
- if not username:
104
- logging.error("Received user_connected signal with no username.")
105
- return jsonify({"error": "Missing 'username' in data"}), 400
106
-
107
- logging.info(f"Received user_connected signal for user: {username}")
108
-
109
- # Call the internal start logic directly (non-blocking? or blocking?)
110
- # Let's make it non-blocking by putting it on the main command queue
111
- # This avoids holding up the signal script and uses the existing processor.
112
- internal_command = {
113
- "action": "_internal_start_ootb",
114
- "target_user": username
115
- }
116
- internal_command_id = f"internal_{username}_{time.time():.0f}"
117
- _service_instance.command_queue.put((internal_command_id, internal_command))
118
- logging.info(f"Queued internal start command {internal_command_id} for user {username}.")
119
-
120
- return jsonify({"message": "Signal received and queued"}), 202
121
-
122
- # --- Helper Functions --- Only logging helpers needed adjustments
123
- # Move these inside the class later
124
- # def get_python_executable(): ...
125
- # def get_pip_executable(): ...
126
-
127
- # Define loggers at module level for use before instance exists?
128
- # Or handle carefully within instance methods.
129
-
130
- # --- PowerShell Task Scheduler Helpers --- (These will become methods) ---
131
-
132
- # _TASK_NAME_PREFIX = "OOTB_UserLogon_" # Move to class
133
-
134
- # def run_powershell_command(command, log_output=True): ...
135
- # def create_or_update_logon_task(username, task_command, python_executable): ...
136
- # def remove_logon_task(username): ...
137
-
138
- # --- End PowerShell Task Scheduler Helpers ---
139
-
140
- TARGET_EXECUTABLE_NAME = "computer-use-ootb-internal.exe"
141
-
142
- class GuardService(win32serviceutil.ServiceFramework):
143
- _svc_name_ = _SERVICE_NAME
144
- _svc_display_name_ = _SERVICE_DISPLAY_NAME
145
- _svc_description_ = _SERVICE_DESCRIPTION
146
- _task_name_prefix = "OOTB_UserLogon_" # Class attribute for task prefix
147
-
148
- # --- Instance Logging Methods ---
149
- def log_info(self, msg):
150
- thread_name = threading.current_thread().name
151
- full_msg = f"[{thread_name}] {msg}"
152
- logging.info(full_msg)
153
- try:
154
- if threading.current_thread().name in ["MainThread", "CommandProcessor"]:
155
- servicemanager.LogInfoMsg(str(full_msg))
156
- except Exception as e:
157
- # Log only to file if event log fails
158
- logging.warning(f"(Instance) Could not write info to Windows Event Log: {e}")
159
-
160
- def log_error(self, msg, exc_info=False):
161
- thread_name = threading.current_thread().name
162
- full_msg = f"[{thread_name}] {msg}"
163
- logging.error(full_msg, exc_info=exc_info)
164
- try:
165
- if threading.current_thread().name in ["MainThread", "CommandProcessor"]:
166
- servicemanager.LogErrorMsg(str(full_msg))
167
- except Exception as e:
168
- logging.warning(f"(Instance) Could not write error to Windows Event Log: {e}")
169
- # --- End Instance Logging ---
170
-
171
- # --- Instance Helper Methods (Moved from module level) ---
172
- def _find_target_executable(self):
173
- """Finds the target executable (e.g., computer-use-ootb-internal.exe) in the Scripts directory."""
174
- try:
175
- # sys.executable should be python.exe or pythonservice.exe in the env root/Scripts
176
- env_dir = os.path.dirname(sys.executable)
177
- # If sys.executable is in Scripts, go up one level
178
- if os.path.basename(env_dir.lower()) == 'scripts':
179
- env_dir = os.path.dirname(env_dir)
180
-
181
- scripts_dir = os.path.join(env_dir, 'Scripts')
182
- target_exe_path = os.path.join(scripts_dir, TARGET_EXECUTABLE_NAME)
183
-
184
- self.log_info(f"_find_target_executable: Checking for executable at: {target_exe_path}")
185
-
186
- if os.path.exists(target_exe_path):
187
- self.log_info(f"_find_target_executable: Found executable: {target_exe_path}")
188
- # Quote if necessary for command line usage
189
- if " " in target_exe_path and not target_exe_path.startswith('"'):
190
- return f'"{target_exe_path}"'
191
- return target_exe_path
192
- else:
193
- self.log_error(f"_find_target_executable: Target executable not found at {target_exe_path}")
194
- # Fallback: Check env root directly (less common for scripts)
195
- target_exe_path_root = os.path.join(env_dir, TARGET_EXECUTABLE_NAME)
196
- self.log_info(f"_find_target_executable: Checking fallback location: {target_exe_path_root}")
197
- if os.path.exists(target_exe_path_root):
198
- self.log_info(f"_find_target_executable: Found executable at fallback location: {target_exe_path_root}")
199
- if " " in target_exe_path_root and not target_exe_path_root.startswith('"'):
200
- return f'"{target_exe_path_root}"'
201
- return target_exe_path_root
202
- else:
203
- self.log_error(f"_find_target_executable: Target executable also not found at {target_exe_path_root}")
204
- return None
205
-
206
- except Exception as e:
207
- self.log_error(f"Error finding target executable: {e}")
208
- return None
209
-
210
- def __init__(self, args):
211
- global _service_instance
212
- win32serviceutil.ServiceFramework.__init__(self, args)
213
- self.hWaitStop = win32event.CreateEvent(None, 0, 0, None)
214
- self.is_running = True
215
- self.server_thread = None
216
- self.command_queue = queue.Queue()
217
- self.command_processor_thread = None
218
- self.session = requests.Session()
219
-
220
- self.target_executable_path = self._find_target_executable()
221
- if not self.target_executable_path:
222
- # Log error and potentially stop service if critical executable is missing
223
- self.log_error(f"CRITICAL: Could not find {TARGET_EXECUTABLE_NAME}. Service cannot function.")
224
- # Consider stopping the service here if needed, or handle appropriately
225
- else:
226
- self.log_info(f"Using target executable: {self.target_executable_path}")
227
-
228
- _service_instance = self
229
- # Determine path to the signal script
230
- self.signal_script_path = self._find_signal_script()
231
- self.log_info(f"Service initialized. Target executable: {self.target_executable_path}. Signal script: {self.signal_script_path}")
232
-
233
- def SvcStop(self):
234
- self.log_info(f"Service stop requested.")
235
- self.ReportServiceStatus(win32service.SERVICE_STOP_PENDING)
236
- self.is_running = False
237
- # Signal the command processor thread to stop
238
- self.command_queue.put(None) # Sentinel value
239
- # Signal the main wait loop
240
- win32event.SetEvent(self.hWaitStop)
241
- # Stopping waitress gracefully from another thread is non-trivial.
242
- # We rely on the SCM timeout / process termination for now.
243
- self.log_info(f"{_SERVICE_NAME} SvcStop: Stop signaled. Server thread will be terminated by SCM.")
244
-
245
-
246
- def SvcDoRun(self):
247
- servicemanager.LogMsg(servicemanager.EVENTLOG_INFORMATION_TYPE,
248
- servicemanager.PYS_SERVICE_STARTED,
249
- (self._svc_name_, ''))
250
- try:
251
- self.log_info(f"{_SERVICE_NAME} starting.")
252
- # Start the command processor thread
253
- self.command_processor_thread = threading.Thread(
254
- target=self.process_commands, name="CommandProcessor", daemon=True)
255
- self.command_processor_thread.start()
256
- self.log_info("Command processor thread started.")
257
-
258
- # Start the Flask server (via Waitress) in a separate thread
259
- self.server_thread = threading.Thread(
260
- target=self.run_server, name="WebServerThread", daemon=True)
261
- self.server_thread.start()
262
- self.log_info(f"Web server thread started, listening on {_LISTEN_HOST}:{_LISTEN_PORT}.")
263
-
264
- self.log_info(f"{_SERVICE_NAME} running. Waiting for stop signal.")
265
- # Keep the main service thread alive waiting for stop signal
266
- win32event.WaitForSingleObject(self.hWaitStop, win32event.INFINITE)
267
- self.log_info(f"{_SERVICE_NAME} received stop signal in main thread.")
268
-
269
- except Exception as e:
270
- self.log_error(f"Fatal error in SvcDoRun: {e}", exc_info=True)
271
- self.SvcStop() # Signal stop if possible
272
- finally:
273
- self.log_info(f"{_SERVICE_NAME} SvcDoRun finished.")
274
-
275
-
276
- def run_server(self):
277
- """Runs the Flask app using Waitress."""
278
- self.log_info(f"Waitress server starting on {_LISTEN_HOST}:{_LISTEN_PORT}")
279
- try:
280
- serve(flask_app, host=_LISTEN_HOST, port=_LISTEN_PORT, threads=4)
281
- self.log_info("Waitress server has stopped.") # Should only happen on shutdown
282
- except Exception as e:
283
- self.log_error(f"Web server thread encountered an error: {e}", exc_info=True)
284
- # Consider signaling the main thread to stop if the web server fails critically
285
- # For now, just log the error.
286
-
287
-
288
- def process_commands(self):
289
- """Worker thread to process commands from the queue."""
290
- self.log_info("Command processor thread starting.")
291
- while self.is_running:
292
- try:
293
- item = self.command_queue.get(block=True, timeout=1) # Add timeout to check is_running periodically
294
- if item is None:
295
- self.log_info("Command processor received stop signal.")
296
- break # Exit loop
297
-
298
- command_id, command = item
299
- action = command.get("action")
300
- target = command.get("target_user", "all_active")
301
- status = "failed_unknown" # Default status for external commands
302
- is_internal = action == "_internal_start_ootb"
303
-
304
- self.log_info(f"Dequeued {'Internal' if is_internal else 'External'} Command ID {command_id}: action='{action}', target='{target}'")
305
-
306
- try:
307
- if action == "update":
308
- status = self.handle_update()
309
- elif action == "stop_ootb":
310
- status = self.handle_stop(target)
311
- elif action == "start_ootb":
312
- status = self.handle_start(target)
313
- elif action == "_internal_start_ootb":
314
- # Call the core start logic but don't report status back externally
315
- internal_status = self._trigger_start_for_user(target)
316
- self.log_info(f"Internal start for {target} finished with status: {internal_status}")
317
- # No external status reporting for internal commands
318
- else:
319
- self.log_error(f"Unknown action in queue: {action}")
320
- status = "failed_unknown_action"
321
- except Exception as handler_ex:
322
- self.log_error(f"Exception processing Command ID {command_id} ({action}): {handler_ex}", exc_info=True)
323
- status = "failed_exception"
324
- finally:
325
- # Only report status for non-internal commands
326
- if not is_internal:
327
- self.report_command_status(command_id, status)
328
- self.command_queue.task_done()
329
-
330
- except queue.Empty:
331
- # Timeout occurred, just loop again and check self.is_running
332
- continue
333
- except Exception as e:
334
- self.log_error(f"Error in command processing loop: {e}", exc_info=True)
335
- if self.is_running:
336
- time.sleep(5)
337
-
338
- self.log_info("Command processor thread finished.")
339
-
340
-
341
- def report_command_status(self, command_id, status, details=""):
342
- """Sends command status back to the server."""
343
- if not _SERVER_STATUS_REPORT_URL:
344
- self.log_error("No server status report URL configured. Skipping report.")
345
- return
346
-
347
- payload = {
348
- "command_id": command_id,
349
- "status": status,
350
- "details": details,
351
- "machine_id": os.getenv('COMPUTERNAME', 'unknown_guard')
352
- }
353
- self.log_info(f"Reporting status for command {command_id}: {status}")
354
- try:
355
- response = self.session.post(_SERVER_STATUS_REPORT_URL, json=payload, timeout=15)
356
- response.raise_for_status()
357
- self.log_info(f"Status report for command {command_id} accepted by server.")
358
- except requests.exceptions.RequestException as e:
359
- self.log_error(f"Failed to report status for command {command_id}: {e}")
360
- except Exception as e:
361
- self.log_error(f"Unexpected error reporting status for command {command_id}: {e}", exc_info=True)
362
-
363
- # --- Command Handlers --- Now call self. for helpers
364
-
365
- def _get_python_executable_from_target_exe(self):
366
- """Attempts to find the python.exe associated with the target executable's env."""
367
- if not self.target_executable_path:
368
- self.log_error("Cannot find python.exe: target executable path is not set.")
369
- return None
370
- try:
371
- exe_path_unquoted = self.target_executable_path.strip('"')
372
- scripts_dir = os.path.dirname(exe_path_unquoted)
373
- # Assume target exe is in a 'Scripts' directory relative to env root
374
- env_dir = os.path.dirname(scripts_dir)
375
- if os.path.basename(scripts_dir.lower()) != 'scripts':
376
- self.log_warning(f"Target executable {exe_path_unquoted} not in expected 'Scripts' directory. Cannot reliably find python.exe.")
377
- # Fallback: maybe the target IS python.exe or next to it?
378
- env_dir = scripts_dir # Try assuming it's env root
379
-
380
- python_exe_path = os.path.join(env_dir, 'python.exe')
381
- self.log_info(f"Checking for python.exe at: {python_exe_path}")
382
- if os.path.exists(python_exe_path):
383
- self.log_info(f"Found associated python.exe: {python_exe_path}")
384
- if " " in python_exe_path and not python_exe_path.startswith('"'):
385
- return f'"{python_exe_path}"'
386
- return python_exe_path
387
- else:
388
- self.log_error(f"Associated python.exe not found at {python_exe_path}")
389
- # Fallback: Check pythonw.exe?
390
- pythonw_exe_path = os.path.join(env_dir, 'pythonw.exe')
391
- if os.path.exists(pythonw_exe_path):
392
- self.log_info(f"Found associated pythonw.exe as fallback: {pythonw_exe_path}")
393
- if " " in pythonw_exe_path and not pythonw_exe_path.startswith('"'):
394
- return f'"{pythonw_exe_path}"'
395
- return pythonw_exe_path
396
- else:
397
- self.log_error(f"Associated pythonw.exe also not found.")
398
- return None
399
-
400
- except Exception as e:
401
- self.log_error(f"Error finding associated python executable: {e}")
402
- return None
403
-
404
- def handle_update(self):
405
- """Handles the update command by running pip install --upgrade directly."""
406
- self.log_info("Executing OOTB update via pip...")
407
-
408
- python_exe = self._get_python_executable_from_target_exe()
409
- if not python_exe:
410
- self.log_error("Cannot update: Could not find associated python.exe for pip.")
411
- return "failed_python_not_found"
412
-
413
- # Package name needs to be defined (replace with actual package name)
414
- package_name = "computer-use-ootb-internal" # Make sure this is correct
415
-
416
- # Construct the command: "C:\path\to\python.exe" -m pip install --upgrade --no-cache-dir package_name
417
- python_exe_unquoted = python_exe.strip('"')
418
- pip_args = ["-m", "pip", "install", "--upgrade", "--no-cache-dir", package_name]
419
- update_command_display = f'{python_exe} {" ".join(pip_args)}'
420
-
421
- self.log_info(f"Running update command: {update_command_display}")
422
- try:
423
- # Execute the pip command directly. Running as LocalSystem should have rights.
424
- result = subprocess.run(
425
- [python_exe_unquoted] + pip_args,
426
- capture_output=True,
427
- text=True,
428
- check=False, # Check manually
429
- encoding='utf-8',
430
- errors='ignore'
431
- )
432
-
433
- if result.stdout:
434
- self.log_info(f"Update process STDOUT:\n{result.stdout.strip()}")
435
- if result.stderr:
436
- self.log_warning(f"Update process STDERR:\n{result.stderr.strip()}")
437
-
438
- if result.returncode == 0:
439
- self.log_info("Update process completed successfully (Exit Code 0).")
440
- return "success"
441
- else:
442
- self.log_error(f"Update process failed (Exit Code {result.returncode}).")
443
- return f"failed_pip_exit_code_{result.returncode}"
444
-
445
- except FileNotFoundError:
446
- self.log_error(f"Update failed: Python executable not found at '{python_exe_unquoted}'.")
447
- return "failed_python_not_found"
448
- except Exception as e:
449
- self.log_error(f"Update failed with exception: {e}", exc_info=True)
450
- return "failed_exception"
451
-
452
- def _get_ootb_processes(self, target_user="all_active"):
453
- ootb_procs = []
454
- target_pid_list = []
455
- try:
456
- target_users = set()
457
- if target_user == "all_active":
458
- for user_session in psutil.users():
459
- username = user_session.name.split('\\')[-1]
460
- target_users.add(username.lower())
461
- else:
462
- target_users.add(target_user.lower())
463
- self.log_info(f"Searching for OOTB processes for users: {target_users}")
464
-
465
- # Use the potentially corrected python.exe path for matching
466
- python_exe_path_for_check = self.target_executable_path.strip('"')
467
- self.log_info(f"_get_ootb_processes: Checking against python path: {python_exe_path_for_check}")
468
-
469
- for proc in psutil.process_iter(['pid', 'name', 'username', 'cmdline', 'exe']):
470
- try:
471
- pinfo = proc.info
472
- proc_username = pinfo['username']
473
- if proc_username:
474
- proc_username = proc_username.split('\\')[-1].lower()
475
-
476
- if proc_username in target_users:
477
- cmdline = ' '.join(pinfo['cmdline']) if pinfo['cmdline'] else ''
478
- # Check if the process executable matches our corrected python path AND module is in cmdline
479
- if pinfo['exe'] and pinfo['exe'] == python_exe_path_for_check and _OOTB_MODULE in cmdline:
480
- self.log_info(f"Found matching OOTB process: PID={pinfo['pid']}, User={pinfo['username']}, Cmd={cmdline}")
481
- ootb_procs.append(proc)
482
- target_pid_list.append(pinfo['pid'])
483
- except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
484
- continue
485
- self.log_info(f"Found {len(ootb_procs)} OOTB process(es) matching criteria: {target_pid_list}")
486
- except Exception as e:
487
- self.log_error(f"Error enumerating processes: {e}", exc_info=True)
488
- return ootb_procs
489
-
490
- def handle_stop(self, target_user="all_active"):
491
- """Stops the OOTB process for specified user(s). Uses psutil first, then taskkill fallback."""
492
- self.log_info(f"Executing stop OOTB for target '{target_user}'...")
493
- stopped_count_psutil = 0
494
- stopped_count_taskkill = 0
495
- errors = []
496
-
497
- target_users_lower = set()
498
- if target_user == "all_active":
499
- try:
500
- for user_session in psutil.users():
501
- username_lower = user_session.name.split('\\')[-1].lower()
502
- target_users_lower.add(username_lower)
503
- self.log_info(f"Targeting all users found by psutil: {target_users_lower}")
504
- except Exception as e:
505
- self.log_error(f"Could not list users via psutil for stop all: {e}")
506
- errors.append("failed_user_enumeration")
507
- target_users_lower = set() # Avoid proceeding if user list failed
508
- else:
509
- target_users_lower.add(target_user.lower())
510
- self.log_info(f"Targeting specific user: {target_user.lower()}")
511
-
512
- if not target_users_lower and target_user == "all_active":
513
- self.log_info("No active users found to stop.")
514
- # If specific user targeted, proceed even if psutil didn't list them (maybe inactive)
515
-
516
- procs_to_kill_by_user = {user: [] for user in target_users_lower}
517
-
518
- # --- Attempt 1: psutil find and terminate ---
519
- self.log_info("Attempting stop using psutil...")
520
- try:
521
- all_running = self._get_ootb_processes("all") # Get all regardless of user first
522
- for proc in all_running:
523
- try:
524
- proc_user = proc.info.get('username')
525
- if proc_user:
526
- user_lower = proc_user.split('\\')[-1].lower()
527
- if user_lower in target_users_lower:
528
- procs_to_kill_by_user[user_lower].append(proc)
529
- except (psutil.NoSuchProcess, psutil.AccessDenied):
530
- pass # Process ended or we can't access it
531
- except Exception as e:
532
- self.log_error(f"Error getting process list for psutil stop: {e}")
533
- errors.append("failed_psutil_list")
534
-
535
- for user, procs in procs_to_kill_by_user.items():
536
- if not procs:
537
- self.log_info(f"psutil: No OOTB processes found running for user '{user}'.")
538
- continue
539
-
540
- self.log_info(f"psutil: Found {len(procs)} OOTB process(es) for user '{user}'. Attempting terminate...")
541
- for proc in procs:
542
- try:
543
- pid = proc.pid
544
- self.log_info(f"psutil: Terminating PID {pid} for user '{user}'...")
545
- proc.terminate()
546
- try:
547
- proc.wait(timeout=3) # Wait a bit for termination
548
- self.log_info(f"psutil: PID {pid} terminated successfully.")
549
- stopped_count_psutil += 1
550
- except psutil.TimeoutExpired:
551
- self.log_warning(f"psutil: PID {pid} did not terminate within timeout. Will try taskkill.")
552
- # No error append yet, let taskkill try
553
- except (psutil.NoSuchProcess, psutil.AccessDenied):
554
- self.log_info(f"psutil: PID {pid} already gone or access denied after terminate.")
555
- stopped_count_psutil += 1 # Count it as stopped
556
- except (psutil.NoSuchProcess, psutil.AccessDenied) as term_err:
557
- self.log_warning(f"psutil: Error terminating PID {proc.pid if 'proc' in locals() and proc else 'unknown'}: {term_err}. It might be gone already.")
558
- # If it's gone, count it?
559
- stopped_count_psutil += 1 # Assume it's gone if NoSuchProcess
560
- except Exception as term_ex:
561
- self.log_error(f"psutil: Unexpected error terminating PID {proc.pid if 'proc' in locals() and proc else 'unknown'}: {term_ex}")
562
- errors.append(f"failed_psutil_terminate_{user}")
563
-
564
- # --- Attempt 2: taskkill fallback (for users where psutil didn't find/stop) ---
565
- # Only run taskkill if psutil didn't stop anything for a specific user OR if target was specific user
566
-
567
- if not self.target_executable_path:
568
- errors.append("skipped_taskkill_no_exe_path")
569
- else:
570
- executable_name = os.path.basename(self.target_executable_path.strip('"'))
571
- for user in target_users_lower:
572
- run_taskkill = False
573
- if user not in procs_to_kill_by_user or not procs_to_kill_by_user[user]:
574
- # psutil didn't find anything for this user initially
575
- run_taskkill = True
576
- self.log_info(f"taskkill: psutil found no processes for '{user}', attempting taskkill as fallback.")
577
- elif any(p.is_running() for p in procs_to_kill_by_user.get(user, []) if p): # Check if any psutil targets still running
578
- run_taskkill = True
579
- self.log_info(f"taskkill: Some processes for '{user}' may remain after psutil, attempting taskkill cleanup.")
580
-
581
- if run_taskkill:
582
- # Construct taskkill command
583
- # Username format might need adjustment (e.g., DOMAIN\user). Try simple first.
584
- taskkill_command = [
585
- "taskkill", "/F", # Force
586
- "/IM", executable_name, # Image name
587
- "/FI", f"USERNAME eq {user}" # Filter by username
588
- ]
589
- self.log_info(f"Running taskkill command: {' '.join(taskkill_command)}")
590
- try:
591
- result = subprocess.run(taskkill_command, capture_output=True, text=True, check=False)
592
- if result.returncode == 0:
593
- self.log_info(f"taskkill successful for user '{user}' (Exit Code 0).")
594
- # Can't easily count how many were killed here, assume success if exit 0
595
- stopped_count_taskkill += 1 # Indicate taskkill ran successfully for user
596
- elif result.returncode == 128: # Code 128: No tasks found matching criteria
597
- self.log_info(f"taskkill: No matching processes found for user '{user}'.")
598
- else:
599
- self.log_error(f"taskkill failed for user '{user}' (Exit Code {result.returncode}).")
600
- self.log_error(f" taskkill STDOUT: {result.stdout.strip()}")
601
- self.log_error(f" taskkill STDERR: {result.stderr.strip()}")
602
- errors.append(f"failed_taskkill_{user}")
603
- except FileNotFoundError:
604
- self.log_error("taskkill command not found.")
605
- errors.append("failed_taskkill_not_found")
606
- break # Stop trying taskkill if command is missing
607
- except Exception as tk_ex:
608
- self.log_error(f"Exception running taskkill for '{user}': {tk_ex}")
609
- errors.append(f"failed_taskkill_exception_{user}")
610
-
611
- # --- Consolidate status ---
612
- final_status = "failed" # Default to failed if errors occurred
613
- if stopped_count_psutil > 0 or stopped_count_taskkill > 0:
614
- final_status = "success" if not errors else "partial_success"
615
- elif not errors:
616
- final_status = "success_no_processes_found"
617
-
618
- details = f"psutil_stopped={stopped_count_psutil}, taskkill_users_attempted={stopped_count_taskkill}, errors={len(errors)}"
619
- self.log_info(f"Finished stopping OOTB. Status: {final_status}. Details: {details}")
620
- return f"{final_status}::{details}" # Return status and details
621
-
622
- def handle_start(self, target_user="all_active"):
623
- """Handles external start command request (finds users, calls internal trigger)."""
624
- self.log_info(f"External start requested for target '{target_user}'...")
625
- # This function now primarily identifies target users and calls the internal trigger method.
626
- # The status returned here reflects the process of identifying and triggering,
627
- # not necessarily the final success/failure of the actual start (which happens async).
628
-
629
- active_sessions = {} # user_lower: session_id
630
- all_system_users = set() # user_lower
631
- try:
632
- # Use psutil for system user list, WTS for active sessions/IDs
633
- for user_session in psutil.users():
634
- username_lower = user_session.name.split('\\')[-1].lower()
635
- all_system_users.add(username_lower)
636
-
637
- sessions = win32ts.WTSEnumerateSessions(win32ts.WTS_CURRENT_SERVER_HANDLE)
638
- for session in sessions:
639
- if session['State'] == win32ts.WTSActive:
640
- try:
641
- user = win32ts.WTSQuerySessionInformation(win32ts.WTS_CURRENT_SERVER_HANDLE, session['SessionId'], win32ts.WTSUserName)
642
- if user:
643
- active_sessions[user.lower()] = session['SessionId']
644
- except Exception as query_err:
645
- self.log_error(f"Could not query session {session['SessionId']}: {query_err}")
646
- except Exception as user_enum_err:
647
- self.log_error(f"Error enumerating users/sessions: {user_enum_err}", exc_info=True)
648
- return "failed_user_enumeration"
649
-
650
- target_users_normalized = set()
651
- if target_user == "all_active":
652
- target_users_normalized = set(active_sessions.keys())
653
- self.log_info(f"Targeting all active users for start: {target_users_normalized}")
654
- else:
655
- normalized_target = target_user.lower()
656
- target_users_normalized.add(normalized_target)
657
- self.log_info(f"Targeting specific user for start: {normalized_target}")
658
-
659
- if not target_users_normalized:
660
- self.log_info("No target users identified for start.")
661
- return "failed_no_target_users"
662
-
663
- trigger_results = {}
664
- for user in target_users_normalized:
665
- self.log_info(f"Calling internal start trigger for user: {user}")
666
- # Call the core logic directly (this is now synchronous within the handler)
667
- # Or queue it? Queuing might be better to avoid blocking the handler if many users.
668
- # Let's stick to the queue approach from the internal endpoint:
669
- internal_command = {
670
- "action": "_internal_start_ootb",
671
- "target_user": user
672
- }
673
- internal_command_id = f"external_{user}_{time.time():.0f}"
674
- self.command_queue.put((internal_command_id, internal_command))
675
- trigger_results[user] = "queued"
676
-
677
- self.log_info(f"Finished queuing start triggers. Results: {trigger_results}")
678
- # The status here just means we successfully queued the actions
679
- # Actual success/failure happens in the command processor later.
680
- # We might need a different way to report overall status if needed.
681
- return f"success_queued::{json.dumps(trigger_results)}"
682
-
683
-
684
- def _trigger_start_for_user(self, username):
685
- """Core logic to start OOTB for a single user. Called internally."""
686
- user = username.lower() # Ensure lowercase
687
- self.log_info(f"Internal trigger: Starting OOTB check for user '{user}'...")
688
- task_created_status = "task_unknown"
689
- immediate_start_status = "start_not_attempted"
690
- final_status = "failed_unknown"
691
-
692
- try:
693
- # 1. Ensure scheduled task exists (still useful fallback/persistence)
694
- try:
695
- task_created = self.create_or_update_logon_task(user)
696
- task_created_status = "task_success" if task_created else "task_failed"
697
- except Exception as task_err:
698
- self.log_error(f"Internal trigger: Exception creating/updating task for {user}: {task_err}", exc_info=True)
699
- task_created_status = "task_exception"
700
- # Don't necessarily fail the whole operation yet
701
-
702
- # 2. Check if user is active
703
- active_sessions = {} # Re-check active sessions specifically for this user
704
- session_id = None
705
- token = None
706
- is_active = False
707
- try:
708
- sessions = win32ts.WTSEnumerateSessions(win32ts.WTS_CURRENT_SERVER_HANDLE)
709
- for session in sessions:
710
- if session['State'] == win32ts.WTSActive:
711
- try:
712
- current_user = win32ts.WTSQuerySessionInformation(win32ts.WTS_CURRENT_SERVER_HANDLE, session['SessionId'], win32ts.WTSUserName)
713
- if current_user and current_user.lower() == user:
714
- session_id = session['SessionId']
715
- is_active = True
716
- self.log_info(f"Internal trigger: User '{user}' is active in session {session_id}.")
717
- break
718
- except Exception: pass # Ignore errors querying other sessions
719
- except Exception as e:
720
- self.log_error(f"Internal trigger: Error checking active sessions for {user}: {e}")
721
- # Continue, assume inactive if check failed?
722
-
723
- if not is_active:
724
- self.log_info(f"Internal trigger: User '{user}' is not active. Skipping immediate start.")
725
- immediate_start_status = "start_skipped_inactive"
726
- final_status = task_created_status # Status depends only on task creation
727
- return final_status # Exit early if inactive
728
-
729
- # 3. Check if already running for this active user
730
- is_running = False
731
- try:
732
- running_procs = self._get_ootb_processes(user)
733
- if running_procs:
734
- is_running = True
735
- self.log_info(f"Internal trigger: OOTB already running for active user '{user}'. Skipping immediate start.")
736
- immediate_start_status = "start_skipped_already_running"
737
- final_status = "success_already_running" # Considered success
738
- return final_status # Exit early if already running
739
- except Exception as e:
740
- self.log_error(f"Internal trigger: Error checking existing processes for {user}: {e}")
741
- # Continue and attempt start despite error?
742
-
743
- # 4. Attempt immediate start (User is active and not running)
744
- immediate_start_status = "start_attempted"
745
- self.log_info(f"Internal trigger: User '{user}' is active and not running. Attempting immediate start via CreateProcessAsUser...")
746
- try:
747
- token = win32ts.WTSQueryUserToken(session_id)
748
- env = win32profile.CreateEnvironmentBlock(token, False)
749
- startup = win32process.STARTUPINFO()
750
- creation_flags = 0x00000010 # CREATE_NEW_CONSOLE
751
- lpApplicationName = None
752
- lpCommandLine = f'cmd.exe /K "{self.target_executable_path}"'
753
- cwd = os.path.dirname(self.target_executable_path.strip('"')) if os.path.dirname(self.target_executable_path.strip('"')) != '' else None
754
-
755
- # Log details before call
756
- self.log_info(f"Internal trigger: Calling CreateProcessAsUser:")
757
- self.log_info(f" lpCommandLine: {lpCommandLine}")
758
- self.log_info(f" lpCurrentDirectory: {cwd if cwd else 'Default'}")
759
-
760
- hProcess, hThread, dwPid, dwTid = win32process.CreateProcessAsUser(
761
- token, lpApplicationName, lpCommandLine, None, None, False,
762
- creation_flags, env, cwd, startup
763
- )
764
- self.log_info(f"Internal trigger: CreateProcessAsUser call succeeded for user '{user}' (PID: {dwPid}). Checking existence...")
765
- win32api.CloseHandle(hProcess)
766
- win32api.CloseHandle(hThread)
767
-
768
- time.sleep(1)
769
- if psutil.pid_exists(dwPid):
770
- self.log_info(f"Internal trigger: Immediate start succeeded for user '{user}' (PID {dwPid}).")
771
- immediate_start_status = "start_success"
772
- final_status = "success" # Overall success
773
- else:
774
- self.log_error(f"Internal trigger: Immediate start failed for user '{user}': Process {dwPid} exited immediately.")
775
- immediate_start_status = "start_failed_exited"
776
- final_status = "failed_start_exited"
777
-
778
- except Exception as proc_err:
779
- self.log_error(f"Internal trigger: Exception during CreateProcessAsUser for user '{user}': {proc_err}", exc_info=True)
780
- immediate_start_status = "start_failed_exception"
781
- final_status = "failed_start_exception"
782
- finally:
783
- if token: win32api.CloseHandle(token)
784
-
785
- # Combine results (though mostly determined by start attempt now)
786
- # Example: final_status = f"{task_created_status}_{immediate_start_status}"
787
- return final_status
788
-
789
- except Exception as e:
790
- self.log_error(f"Internal trigger: Unexpected error processing start for {username}: {e}", exc_info=True)
791
- return "failed_trigger_exception"
792
-
793
-
794
- def create_or_update_logon_task(self, username):
795
- """Creates/updates task to trigger the internal signal script on session connect."""
796
- if not self.signal_script_path:
797
- self.log_error(f"Cannot create task for {username}: Signal script path is not set.")
798
- return False
799
- if not sys.executable:
800
- self.log_error(f"Cannot create task for {username}: sys.executable is not found.")
801
- return False
802
-
803
- # Use the python executable that the service itself is running under
804
- python_exe = sys.executable
805
- if ' ' in python_exe and not python_exe.startswith('"'):
806
- python_exe = f'"{python_exe}"'
807
-
808
- task_name = f"OOTB_UserConnect_{username}"
809
- # Action: Run python.exe with the signal script and username argument
810
- action_executable = python_exe
811
- # Ensure script path is quoted if needed
812
- script_arg = self.signal_script_path # Should be quoted already by _find_signal_script
813
- # Username might need quoting if it contains spaces, though unlikely
814
- user_arg = username # Keep simple for now
815
- action_arguments = f'{script_arg} "{user_arg}"' # Pass username as quoted arg
816
- safe_action_executable = action_executable.replace("'", "''") # Escape for PS
817
- safe_action_arguments = action_arguments.replace("'", "''") # Escape for PS
818
-
819
- # Working directory for the script (likely its own directory)
820
- try:
821
- script_dir = os.path.dirname(self.signal_script_path.strip('"'))
822
- if not script_dir: script_dir = "."
823
- safe_working_directory = script_dir.replace("'", "''")
824
- working_directory_setting = f"$action.WorkingDirectory = '{safe_working_directory}'"
825
- except Exception as e:
826
- self.log_error(f"Error determining working directory for signal script task: {e}. WD will not be set.")
827
- working_directory_setting = "# Could not set WorkingDirectory"
828
-
829
- # PowerShell command construction
830
- ps_command = f"""
831
- $taskName = "{task_name}"
832
- $principal = New-ScheduledTaskPrincipal -UserId "{username}" -LogonType Interactive
833
-
834
- # Action: Run python signal script
835
- $action = New-ScheduledTaskAction -Execute '{safe_action_executable}' -Argument '{safe_action_arguments}'
836
- {working_directory_setting}
837
-
838
- # Trigger: On session connect (Event ID 21)
839
- $logName = 'Microsoft-Windows-TerminalServices-LocalSessionManager/Operational'
840
- $source = 'Microsoft-Windows-TerminalServices-LocalSessionManager'
841
- $eventIDs = @(21, 25)
842
- $trigger = New-ScheduledTaskTrigger -Event -LogName $logName -Source $source -EventId $eventIDs[0]
843
- # Optional Delay: -Delay 'PT15S'
844
- # $trigger.Delay = 'PT15S'
845
-
846
- $settings = New-ScheduledTaskSettingsSet -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries -ExecutionTimeLimit (New-TimeSpan -Days 9999) -RestartCount 3 -RestartInterval (New-TimeSpan -Minutes 1)
847
- $description = "Triggers OOTB Guard Service for user {username} upon session connect via internal signal." # Updated description
848
-
849
- # Unregister existing task first (force)
850
- Unregister-ScheduledTask -TaskName $taskName -Confirm:$false -ErrorAction SilentlyContinue
851
-
852
- # Register the new task
853
- Register-ScheduledTask -TaskName $taskName -Principal $principal -Action $action -Trigger $trigger -Settings $settings -Description $description -Force
854
- """
855
- self.log_info(f"Attempting to create/update task '{task_name}' for user '{username}' to run signal script.")
856
- try:
857
- success = self.run_powershell_command(ps_command)
858
- if success:
859
- self.log_info(f"Successfully ran PowerShell command to create/update task '{task_name}'.")
860
- return True
861
- else:
862
- self.log_error(f"PowerShell command failed to create/update task '{task_name}'. See previous logs.")
863
- return False
864
- except Exception as e:
865
- self.log_error(f"Failed to create/update scheduled task '{task_name}' for user '{username}': {e}", exc_info=True)
866
- return False
867
-
868
- def run_powershell_command(self, command, log_output=True):
869
- """Executes a PowerShell command and handles output/errors. Returns True on success."""
870
- self.log_info(f"Executing PowerShell: {command}")
871
- try:
872
- result = subprocess.run(
873
- ["powershell.exe", "-NoProfile", "-NonInteractive", "-Command", command],
874
- capture_output=True, text=True, check=True, encoding='utf-8', errors='ignore'
875
- )
876
- if log_output and result.stdout:
877
- self.log_info(f"PowerShell STDOUT:\n{result.stdout.strip()}")
878
- if log_output and result.stderr:
879
- # Log stderr as info, as some commands write status here (like unregister task not found)
880
- self.log_info(f"PowerShell STDERR:\n{result.stderr.strip()}")
881
- return True
882
- except FileNotFoundError:
883
- self.log_error("'powershell.exe' not found. Cannot manage scheduled tasks.")
884
- return False
885
- except subprocess.CalledProcessError as e:
886
- # Log error but still return False, handled by caller
887
- self.log_error(f"PowerShell command failed (Exit Code {e.returncode}):")
888
- self.log_error(f" Command: {e.cmd}")
889
- if e.stdout: self.log_error(f" STDOUT: {e.stdout.strip()}")
890
- if e.stderr: self.log_error(f" STDERR: {e.stderr.strip()}")
891
- return False
892
- except Exception as e:
893
- self.log_error(f"Unexpected error running PowerShell: {e}", exc_info=True)
894
- return False
895
-
896
- def remove_logon_task(self, username):
897
- """Removes the logon scheduled task for a user."""
898
- task_name = f"{self._task_name_prefix}{username}"
899
- safe_task_name = task_name.replace("'", "''")
900
- command = f"Unregister-ScheduledTask -TaskName '{safe_task_name}' -Confirm:$false -ErrorAction SilentlyContinue"
901
- self.run_powershell_command(command, log_output=False)
902
- self.log_info(f"Attempted removal of scheduled task '{task_name}' for user '{username}'.")
903
- return True
904
-
905
- def _find_signal_script(self):
906
- """Finds the signal_connection.py script relative to this service file."""
907
- try:
908
- base_dir = os.path.dirname(os.path.abspath(__file__))
909
- script_path = os.path.join(base_dir, "signal_connection.py")
910
- if os.path.exists(script_path):
911
- self.log_info(f"Found signal script at: {script_path}")
912
- # Quote if needed?
913
- if " " in script_path and not script_path.startswith('"'):
914
- return f'"{script_path}"'
915
- return script_path
916
- else:
917
- self.log_error(f"Signal script signal_connection.py not found near {base_dir}")
918
- return None
919
- except Exception as e:
920
- self.log_error(f"Error finding signal script: {e}")
921
- return None
922
-
923
- # --- Main Execution Block ---
924
- if __name__ == '__main__':
925
- if len(sys.argv) > 1 and sys.argv[1] == 'debug':
926
- self.log_info("Starting service in debug mode...")
927
- print(f"Running Flask server via Waitress on {_LISTEN_HOST}:{_LISTEN_PORT} for debugging...")
928
- print("Service logic (command processing) will NOT run in this mode.")
929
- print("Use this primarily to test the '/command' endpoint receiving POSTs.")
930
- print("Press Ctrl+C to stop.")
931
- try:
932
- serve(flask_app, host=_LISTEN_HOST, port=_LISTEN_PORT, threads=1)
933
- except KeyboardInterrupt:
934
- print("\nDebug server stopped.")
935
-
936
- elif len(sys.argv) == 1:
937
- try:
938
- servicemanager.Initialize()
939
- servicemanager.PrepareToHostSingle(GuardService)
940
- servicemanager.StartServiceCtrlDispatcher()
941
- except win32service.error as details:
942
- import winerror
943
- if details.winerror == winerror.ERROR_FAILED_SERVICE_CONTROLLER_CONNECT:
944
- print(f"Error: Not started by SCM.")
945
- print(f"Use 'python {os.path.basename(__file__)} install|start|stop|remove|debug'")
946
- else:
947
- print(f"Error preparing service: {details}")
948
- except Exception as e:
949
- print(f"Unexpected error initializing service: {e}")
950
- else:
1
+ # src/computer_use_ootb_internal/guard_service.py
2
+ import sys
3
+ import os
4
+ import time
5
+ import logging
6
+ import subprocess
7
+ import pathlib
8
+ import ctypes
9
+ import threading # For running server thread
10
+ import queue # For queuing commands
11
+ import requests # Keep for status reporting back
12
+ import servicemanager # From pywin32
13
+ import win32serviceutil # From pywin32
14
+ import win32service # From pywin32
15
+ import win32event # From pywin32
16
+ import win32api # From pywin32
17
+ import win32process # From pywin32
18
+ import win32security # From pywin32
19
+ import win32profile # From pywin32
20
+ import win32ts # From pywin32 (Terminal Services API)
21
+ import win32con
22
+ import psutil # For process/user info
23
+ from flask import Flask, request, jsonify # For embedded server
24
+ from waitress import serve # For serving Flask app
25
+ import json # Needed for status reporting
26
+
27
+ # --- Configuration ---
28
+ _SERVICE_NAME = "OOTBGuardService"
29
+ _SERVICE_DISPLAY_NAME = "OOTB Guard Service"
30
+ _SERVICE_DESCRIPTION = "Background service for OOTB monitoring and remote management (Server POST mode)."
31
+ _PACKAGE_NAME = "computer-use-ootb-internal"
32
+ _OOTB_MODULE = "computer_use_ootb_internal.app_teachmode"
33
+ # --- Server POST Configuration ---
34
+ _LISTEN_HOST = "0.0.0.0" # Listen on all interfaces
35
+ _LISTEN_PORT = 14000 # Port for server to POST commands TO
36
+ # _SHARED_SECRET = "YOUR_SECRET_HERE" # !! REMOVED !! - No secret check implemented now
37
+ # --- End Server POST Configuration ---
38
+ _SERVER_STATUS_REPORT_URL = "http://52.160.105.102:7000/guard/status" # URL to POST status back TO (Path changed)
39
+ _LOG_FILE = pathlib.Path(os.environ['PROGRAMDATA']) / "OOTBGuardService" / "guard_post_mode.log" # Different log file
40
+ # --- End Configuration ---
41
+
42
+ _LOG_FILE.parent.mkdir(parents=True, exist_ok=True)
43
+ logging.basicConfig(
44
+ filename=_LOG_FILE,
45
+ level=logging.INFO,
46
+ format='%(asctime)s %(levelname)s:%(name)s:%(threadName)s: %(message)s'
47
+ )
48
+
49
+ # --- Global service instance reference (needed for Flask routes) ---
50
+ _service_instance = None
51
+
52
+ # --- Flask App Definition ---
53
+ flask_app = Flask(__name__)
54
+
55
+ @flask_app.route('/command', methods=['POST'])
56
+ def receive_command():
57
+ global _service_instance
58
+ if not _service_instance:
59
+ logging.error("Received command but service instance is not set.")
60
+ return jsonify({"error": "Service not ready"}), 503
61
+
62
+ # --- Authentication REMOVED ---
63
+ # secret = request.headers.get('X-Guard-Secret')
64
+ # if not secret or secret != _SHARED_SECRET:
65
+ # logging.warning(f"Unauthorized command POST received (Invalid/Missing X-Guard-Secret). Remote Addr: {request.remote_addr}")
66
+ # return jsonify({"error": "Unauthorized"}), 403
67
+ # --- End Authentication REMOVED ---
68
+
69
+ if not request.is_json:
70
+ logging.warning("Received non-JSON command POST.")
71
+ return jsonify({"error": "Request must be JSON"}), 400
72
+
73
+ command = request.get_json()
74
+ logging.info(f"Received command via POST: {command}")
75
+
76
+ # Basic validation
77
+ action = command.get("action")
78
+ command_id = command.get("command_id", "N/A") # Use for status reporting
79
+ if not action:
80
+ logging.error(f"Received command POST with no action: {command}")
81
+ return jsonify({"error": "Missing 'action' in command"}), 400
82
+
83
+ # Queue the command for processing in a background thread
84
+ _service_instance.command_queue.put((command_id, command))
85
+ logging.info(f"Queued command {command_id} ({action}) for processing.")
86
+
87
+ return jsonify({"message": f"Command {command_id} received and queued"}), 202 # Accepted
88
+
89
+ @flask_app.route('/internal/user_connected', methods=['POST'])
90
+ def user_connected_notification():
91
+ """Internal endpoint triggered by the scheduled task helper script."""
92
+ global _service_instance
93
+ if not _service_instance:
94
+ logging.error("Received user_connected signal but service instance is not set.")
95
+ return jsonify({"error": "Service not ready"}), 503
96
+
97
+ if not request.is_json:
98
+ logging.warning("Received non-JSON user_connected signal.")
99
+ return jsonify({"error": "Request must be JSON"}), 400
100
+
101
+ data = request.get_json()
102
+ username = data.get("username")
103
+ if not username:
104
+ logging.error("Received user_connected signal with no username.")
105
+ return jsonify({"error": "Missing 'username' in data"}), 400
106
+
107
+ logging.info(f"Received user_connected signal for user: {username}")
108
+
109
+ # Call the internal start logic directly (non-blocking? or blocking?)
110
+ # Let's make it non-blocking by putting it on the main command queue
111
+ # This avoids holding up the signal script and uses the existing processor.
112
+ internal_command = {
113
+ "action": "_internal_start_ootb",
114
+ "target_user": username
115
+ }
116
+ internal_command_id = f"internal_{username}_{time.time():.0f}"
117
+ _service_instance.command_queue.put((internal_command_id, internal_command))
118
+ logging.info(f"Queued internal start command {internal_command_id} for user {username}.")
119
+
120
+ return jsonify({"message": "Signal received and queued"}), 202
121
+
122
+ # --- Helper Functions --- Only logging helpers needed adjustments
123
+ # Move these inside the class later
124
+ # def get_python_executable(): ...
125
+ # def get_pip_executable(): ...
126
+
127
+ # Define loggers at module level for use before instance exists?
128
+ # Or handle carefully within instance methods.
129
+
130
+ # --- PowerShell Task Scheduler Helpers --- (These will become methods) ---
131
+
132
+ # _TASK_NAME_PREFIX = "OOTB_UserLogon_" # Move to class
133
+
134
+ # def run_powershell_command(command, log_output=True): ...
135
+ # def create_or_update_logon_task(username, task_command, python_executable): ...
136
+ # def remove_logon_task(username): ...
137
+
138
+ # --- End PowerShell Task Scheduler Helpers ---
139
+
140
+ TARGET_EXECUTABLE_NAME = "computer-use-ootb-internal.exe"
141
+
142
+ class GuardService(win32serviceutil.ServiceFramework):
143
+ _svc_name_ = _SERVICE_NAME
144
+ _svc_display_name_ = _SERVICE_DISPLAY_NAME
145
+ _svc_description_ = _SERVICE_DESCRIPTION
146
+ _task_name_prefix = "OOTB_UserLogon_" # Class attribute for task prefix
147
+
148
+ # --- Instance Logging Methods ---
149
+ def log_info(self, msg):
150
+ thread_name = threading.current_thread().name
151
+ full_msg = f"[{thread_name}] {msg}"
152
+ logging.info(full_msg)
153
+ try:
154
+ if threading.current_thread().name in ["MainThread", "CommandProcessor"]:
155
+ servicemanager.LogInfoMsg(str(full_msg))
156
+ except Exception as e:
157
+ # Log only to file if event log fails
158
+ logging.warning(f"(Instance) Could not write info to Windows Event Log: {e}")
159
+
160
+ def log_error(self, msg, exc_info=False):
161
+ thread_name = threading.current_thread().name
162
+ full_msg = f"[{thread_name}] {msg}"
163
+ logging.error(full_msg, exc_info=exc_info)
164
+ try:
165
+ if threading.current_thread().name in ["MainThread", "CommandProcessor"]:
166
+ servicemanager.LogErrorMsg(str(full_msg))
167
+ except Exception as e:
168
+ logging.warning(f"(Instance) Could not write error to Windows Event Log: {e}")
169
+ # --- End Instance Logging ---
170
+
171
+ # --- Instance Helper Methods (Moved from module level) ---
172
+ def _find_target_executable(self):
173
+ """Finds the target executable (e.g., computer-use-ootb-internal.exe) in the Scripts directory."""
174
+ try:
175
+ # sys.executable should be python.exe or pythonservice.exe in the env root/Scripts
176
+ env_dir = os.path.dirname(sys.executable)
177
+ # If sys.executable is in Scripts, go up one level
178
+ if os.path.basename(env_dir.lower()) == 'scripts':
179
+ env_dir = os.path.dirname(env_dir)
180
+
181
+ scripts_dir = os.path.join(env_dir, 'Scripts')
182
+ target_exe_path = os.path.join(scripts_dir, TARGET_EXECUTABLE_NAME)
183
+
184
+ self.log_info(f"_find_target_executable: Checking for executable at: {target_exe_path}")
185
+
186
+ if os.path.exists(target_exe_path):
187
+ self.log_info(f"_find_target_executable: Found executable: {target_exe_path}")
188
+ # Quote if necessary for command line usage
189
+ if " " in target_exe_path and not target_exe_path.startswith('"'):
190
+ return f'"{target_exe_path}"'
191
+ return target_exe_path
192
+ else:
193
+ self.log_error(f"_find_target_executable: Target executable not found at {target_exe_path}")
194
+ # Fallback: Check env root directly (less common for scripts)
195
+ target_exe_path_root = os.path.join(env_dir, TARGET_EXECUTABLE_NAME)
196
+ self.log_info(f"_find_target_executable: Checking fallback location: {target_exe_path_root}")
197
+ if os.path.exists(target_exe_path_root):
198
+ self.log_info(f"_find_target_executable: Found executable at fallback location: {target_exe_path_root}")
199
+ if " " in target_exe_path_root and not target_exe_path_root.startswith('"'):
200
+ return f'"{target_exe_path_root}"'
201
+ return target_exe_path_root
202
+ else:
203
+ self.log_error(f"_find_target_executable: Target executable also not found at {target_exe_path_root}")
204
+ return None
205
+
206
+ except Exception as e:
207
+ self.log_error(f"Error finding target executable: {e}")
208
+ return None
209
+
210
+ def __init__(self, args):
211
+ global _service_instance
212
+ win32serviceutil.ServiceFramework.__init__(self, args)
213
+ self.hWaitStop = win32event.CreateEvent(None, 0, 0, None)
214
+ self.is_running = True
215
+ self.server_thread = None
216
+ self.command_queue = queue.Queue()
217
+ self.command_processor_thread = None
218
+ self.session = requests.Session()
219
+
220
+ self.target_executable_path = self._find_target_executable()
221
+ if not self.target_executable_path:
222
+ # Log error and potentially stop service if critical executable is missing
223
+ self.log_error(f"CRITICAL: Could not find {TARGET_EXECUTABLE_NAME}. Service cannot function.")
224
+ # Consider stopping the service here if needed, or handle appropriately
225
+ else:
226
+ self.log_info(f"Using target executable: {self.target_executable_path}")
227
+
228
+ _service_instance = self
229
+ # Determine path to the signal script
230
+ self.signal_script_path = self._find_signal_script()
231
+ self.log_info(f"Service initialized. Target executable: {self.target_executable_path}. Signal script: {self.signal_script_path}")
232
+
233
+ def SvcStop(self):
234
+ self.log_info(f"Service stop requested.")
235
+ self.ReportServiceStatus(win32service.SERVICE_STOP_PENDING)
236
+ self.is_running = False
237
+ # Signal the command processor thread to stop
238
+ self.command_queue.put(None) # Sentinel value
239
+ # Signal the main wait loop
240
+ win32event.SetEvent(self.hWaitStop)
241
+ # Stopping waitress gracefully from another thread is non-trivial.
242
+ # We rely on the SCM timeout / process termination for now.
243
+ self.log_info(f"{_SERVICE_NAME} SvcStop: Stop signaled. Server thread will be terminated by SCM.")
244
+
245
+
246
+ def SvcDoRun(self):
247
+ servicemanager.LogMsg(servicemanager.EVENTLOG_INFORMATION_TYPE,
248
+ servicemanager.PYS_SERVICE_STARTED,
249
+ (self._svc_name_, ''))
250
+ try:
251
+ self.log_info(f"{_SERVICE_NAME} starting.")
252
+ # Start the command processor thread
253
+ self.command_processor_thread = threading.Thread(
254
+ target=self.process_commands, name="CommandProcessor", daemon=True)
255
+ self.command_processor_thread.start()
256
+ self.log_info("Command processor thread started.")
257
+
258
+ # Start the Flask server (via Waitress) in a separate thread
259
+ self.server_thread = threading.Thread(
260
+ target=self.run_server, name="WebServerThread", daemon=True)
261
+ self.server_thread.start()
262
+ self.log_info(f"Web server thread started, listening on {_LISTEN_HOST}:{_LISTEN_PORT}.")
263
+
264
+ self.log_info(f"{_SERVICE_NAME} running. Waiting for stop signal.")
265
+ # Keep the main service thread alive waiting for stop signal
266
+ win32event.WaitForSingleObject(self.hWaitStop, win32event.INFINITE)
267
+ self.log_info(f"{_SERVICE_NAME} received stop signal in main thread.")
268
+
269
+ except Exception as e:
270
+ self.log_error(f"Fatal error in SvcDoRun: {e}", exc_info=True)
271
+ self.SvcStop() # Signal stop if possible
272
+ finally:
273
+ self.log_info(f"{_SERVICE_NAME} SvcDoRun finished.")
274
+
275
+
276
+ def run_server(self):
277
+ """Runs the Flask app using Waitress."""
278
+ self.log_info(f"Waitress server starting on {_LISTEN_HOST}:{_LISTEN_PORT}")
279
+ try:
280
+ serve(flask_app, host=_LISTEN_HOST, port=_LISTEN_PORT, threads=4)
281
+ self.log_info("Waitress server has stopped.") # Should only happen on shutdown
282
+ except Exception as e:
283
+ self.log_error(f"Web server thread encountered an error: {e}", exc_info=True)
284
+ # Consider signaling the main thread to stop if the web server fails critically
285
+ # For now, just log the error.
286
+
287
+
288
+ def process_commands(self):
289
+ """Worker thread to process commands from the queue."""
290
+ self.log_info("Command processor thread starting.")
291
+ while self.is_running:
292
+ try:
293
+ item = self.command_queue.get(block=True, timeout=1) # Add timeout to check is_running periodically
294
+ if item is None:
295
+ self.log_info("Command processor received stop signal.")
296
+ break # Exit loop
297
+
298
+ command_id, command = item
299
+ action = command.get("action")
300
+ target = command.get("target_user", "all_active")
301
+ status = "failed_unknown" # Default status for external commands
302
+ is_internal = action == "_internal_start_ootb"
303
+
304
+ self.log_info(f"Dequeued {'Internal' if is_internal else 'External'} Command ID {command_id}: action='{action}', target='{target}'")
305
+
306
+ try:
307
+ if action == "update":
308
+ status = self.handle_update()
309
+ elif action == "stop_ootb":
310
+ status = self.handle_stop(target)
311
+ elif action == "start_ootb":
312
+ status = self.handle_start(target)
313
+ elif action == "_internal_start_ootb":
314
+ # Call the core start logic but don't report status back externally
315
+ internal_status = self._trigger_start_for_user(target)
316
+ self.log_info(f"Internal start for {target} finished with status: {internal_status}")
317
+ # No external status reporting for internal commands
318
+ else:
319
+ self.log_error(f"Unknown action in queue: {action}")
320
+ status = "failed_unknown_action"
321
+ except Exception as handler_ex:
322
+ self.log_error(f"Exception processing Command ID {command_id} ({action}): {handler_ex}", exc_info=True)
323
+ status = "failed_exception"
324
+ finally:
325
+ # Only report status for non-internal commands
326
+ if not is_internal:
327
+ self.report_command_status(command_id, status)
328
+ self.command_queue.task_done()
329
+
330
+ except queue.Empty:
331
+ # Timeout occurred, just loop again and check self.is_running
332
+ continue
333
+ except Exception as e:
334
+ self.log_error(f"Error in command processing loop: {e}", exc_info=True)
335
+ if self.is_running:
336
+ time.sleep(5)
337
+
338
+ self.log_info("Command processor thread finished.")
339
+
340
+
341
+ def report_command_status(self, command_id, status, details=""):
342
+ """Sends command status back to the server."""
343
+ if not _SERVER_STATUS_REPORT_URL:
344
+ self.log_error("No server status report URL configured. Skipping report.")
345
+ return
346
+
347
+ payload = {
348
+ "command_id": command_id,
349
+ "status": status,
350
+ "details": details,
351
+ "machine_id": os.getenv('COMPUTERNAME', 'unknown_guard')
352
+ }
353
+ self.log_info(f"Reporting status for command {command_id}: {status}")
354
+ try:
355
+ response = self.session.post(_SERVER_STATUS_REPORT_URL, json=payload, timeout=15)
356
+ response.raise_for_status()
357
+ self.log_info(f"Status report for command {command_id} accepted by server.")
358
+ except requests.exceptions.RequestException as e:
359
+ self.log_error(f"Failed to report status for command {command_id}: {e}")
360
+ except Exception as e:
361
+ self.log_error(f"Unexpected error reporting status for command {command_id}: {e}", exc_info=True)
362
+
363
+ # --- Command Handlers --- Now call self. for helpers
364
+
365
+ def _get_python_executable_from_target_exe(self):
366
+ """Attempts to find the python.exe associated with the target executable's env."""
367
+ if not self.target_executable_path:
368
+ self.log_error("Cannot find python.exe: target executable path is not set.")
369
+ return None
370
+ try:
371
+ exe_path_unquoted = self.target_executable_path.strip('"')
372
+ scripts_dir = os.path.dirname(exe_path_unquoted)
373
+ # Assume target exe is in a 'Scripts' directory relative to env root
374
+ env_dir = os.path.dirname(scripts_dir)
375
+ if os.path.basename(scripts_dir.lower()) != 'scripts':
376
+ self.log_warning(f"Target executable {exe_path_unquoted} not in expected 'Scripts' directory. Cannot reliably find python.exe.")
377
+ # Fallback: maybe the target IS python.exe or next to it?
378
+ env_dir = scripts_dir # Try assuming it's env root
379
+
380
+ python_exe_path = os.path.join(env_dir, 'python.exe')
381
+ self.log_info(f"Checking for python.exe at: {python_exe_path}")
382
+ if os.path.exists(python_exe_path):
383
+ self.log_info(f"Found associated python.exe: {python_exe_path}")
384
+ if " " in python_exe_path and not python_exe_path.startswith('"'):
385
+ return f'"{python_exe_path}"'
386
+ return python_exe_path
387
+ else:
388
+ self.log_error(f"Associated python.exe not found at {python_exe_path}")
389
+ # Fallback: Check pythonw.exe?
390
+ pythonw_exe_path = os.path.join(env_dir, 'pythonw.exe')
391
+ if os.path.exists(pythonw_exe_path):
392
+ self.log_info(f"Found associated pythonw.exe as fallback: {pythonw_exe_path}")
393
+ if " " in pythonw_exe_path and not pythonw_exe_path.startswith('"'):
394
+ return f'"{pythonw_exe_path}"'
395
+ return pythonw_exe_path
396
+ else:
397
+ self.log_error(f"Associated pythonw.exe also not found.")
398
+ return None
399
+
400
+ except Exception as e:
401
+ self.log_error(f"Error finding associated python executable: {e}")
402
+ return None
403
+
404
+ def handle_update(self):
405
+ """Handles the update command by running pip install --upgrade directly."""
406
+ self.log_info("Executing OOTB update via pip...")
407
+
408
+ python_exe = self._get_python_executable_from_target_exe()
409
+ if not python_exe:
410
+ self.log_error("Cannot update: Could not find associated python.exe for pip.")
411
+ return "failed_python_not_found"
412
+
413
+ # Package name needs to be defined (replace with actual package name)
414
+ package_name = "computer-use-ootb-internal" # Make sure this is correct
415
+
416
+ # Construct the command: "C:\path\to\python.exe" -m pip install --upgrade --no-cache-dir package_name
417
+ python_exe_unquoted = python_exe.strip('"')
418
+ pip_args = ["-m", "pip", "install", "--upgrade", "--no-cache-dir", package_name]
419
+ update_command_display = f'{python_exe} {" ".join(pip_args)}'
420
+
421
+ self.log_info(f"Running update command: {update_command_display}")
422
+ try:
423
+ # Execute the pip command directly. Running as LocalSystem should have rights.
424
+ result = subprocess.run(
425
+ [python_exe_unquoted] + pip_args,
426
+ capture_output=True,
427
+ text=True,
428
+ check=False, # Check manually
429
+ encoding='utf-8',
430
+ errors='ignore'
431
+ )
432
+
433
+ if result.stdout:
434
+ self.log_info(f"Update process STDOUT:\n{result.stdout.strip()}")
435
+ if result.stderr:
436
+ self.log_warning(f"Update process STDERR:\n{result.stderr.strip()}")
437
+
438
+ if result.returncode == 0:
439
+ self.log_info("Update process completed successfully (Exit Code 0).")
440
+ return "success"
441
+ else:
442
+ self.log_error(f"Update process failed (Exit Code {result.returncode}).")
443
+ return f"failed_pip_exit_code_{result.returncode}"
444
+
445
+ except FileNotFoundError:
446
+ self.log_error(f"Update failed: Python executable not found at '{python_exe_unquoted}'.")
447
+ return "failed_python_not_found"
448
+ except Exception as e:
449
+ self.log_error(f"Update failed with exception: {e}", exc_info=True)
450
+ return "failed_exception"
451
+
452
+ def _get_ootb_processes(self, target_user="all_active"):
453
+ ootb_procs = []
454
+ target_pid_list = []
455
+ try:
456
+ target_users = set()
457
+ if target_user == "all_active":
458
+ for user_session in psutil.users():
459
+ username = user_session.name.split('\\')[-1]
460
+ target_users.add(username.lower())
461
+ else:
462
+ target_users.add(target_user.lower())
463
+ self.log_info(f"Searching for OOTB processes for users: {target_users}")
464
+
465
+ # Use the potentially corrected python.exe path for matching
466
+ python_exe_path_for_check = self.target_executable_path.strip('"')
467
+ self.log_info(f"_get_ootb_processes: Checking against python path: {python_exe_path_for_check}")
468
+
469
+ for proc in psutil.process_iter(['pid', 'name', 'username', 'cmdline', 'exe']):
470
+ try:
471
+ pinfo = proc.info
472
+ proc_username = pinfo['username']
473
+ if proc_username:
474
+ proc_username = proc_username.split('\\')[-1].lower()
475
+
476
+ if proc_username in target_users:
477
+ cmdline = ' '.join(pinfo['cmdline']) if pinfo['cmdline'] else ''
478
+ # Check if the process executable matches our corrected python path AND module is in cmdline
479
+ if pinfo['exe'] and pinfo['exe'] == python_exe_path_for_check and _OOTB_MODULE in cmdline:
480
+ self.log_info(f"Found matching OOTB process: PID={pinfo['pid']}, User={pinfo['username']}, Cmd={cmdline}")
481
+ ootb_procs.append(proc)
482
+ target_pid_list.append(pinfo['pid'])
483
+ except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
484
+ continue
485
+ self.log_info(f"Found {len(ootb_procs)} OOTB process(es) matching criteria: {target_pid_list}")
486
+ except Exception as e:
487
+ self.log_error(f"Error enumerating processes: {e}", exc_info=True)
488
+ return ootb_procs
489
+
490
+ def handle_stop(self, target_user="all_active"):
491
+ """Stops the OOTB process for specified user(s). Uses psutil first, then taskkill fallback."""
492
+ self.log_info(f"Executing stop OOTB for target '{target_user}'...")
493
+ stopped_count_psutil = 0
494
+ stopped_count_taskkill = 0
495
+ errors = []
496
+
497
+ target_users_lower = set()
498
+ if target_user == "all_active":
499
+ try:
500
+ for user_session in psutil.users():
501
+ username_lower = user_session.name.split('\\')[-1].lower()
502
+ target_users_lower.add(username_lower)
503
+ self.log_info(f"Targeting all users found by psutil: {target_users_lower}")
504
+ except Exception as e:
505
+ self.log_error(f"Could not list users via psutil for stop all: {e}")
506
+ errors.append("failed_user_enumeration")
507
+ target_users_lower = set() # Avoid proceeding if user list failed
508
+ else:
509
+ target_users_lower.add(target_user.lower())
510
+ self.log_info(f"Targeting specific user: {target_user.lower()}")
511
+
512
+ if not target_users_lower and target_user == "all_active":
513
+ self.log_info("No active users found to stop.")
514
+ # If specific user targeted, proceed even if psutil didn't list them (maybe inactive)
515
+
516
+ procs_to_kill_by_user = {user: [] for user in target_users_lower}
517
+
518
+ # --- Attempt 1: psutil find and terminate ---
519
+ self.log_info("Attempting stop using psutil...")
520
+ try:
521
+ all_running = self._get_ootb_processes("all") # Get all regardless of user first
522
+ for proc in all_running:
523
+ try:
524
+ proc_user = proc.info.get('username')
525
+ if proc_user:
526
+ user_lower = proc_user.split('\\')[-1].lower()
527
+ if user_lower in target_users_lower:
528
+ procs_to_kill_by_user[user_lower].append(proc)
529
+ except (psutil.NoSuchProcess, psutil.AccessDenied):
530
+ pass # Process ended or we can't access it
531
+ except Exception as e:
532
+ self.log_error(f"Error getting process list for psutil stop: {e}")
533
+ errors.append("failed_psutil_list")
534
+
535
+ for user, procs in procs_to_kill_by_user.items():
536
+ if not procs:
537
+ self.log_info(f"psutil: No OOTB processes found running for user '{user}'.")
538
+ continue
539
+
540
+ self.log_info(f"psutil: Found {len(procs)} OOTB process(es) for user '{user}'. Attempting terminate...")
541
+ for proc in procs:
542
+ try:
543
+ pid = proc.pid
544
+ self.log_info(f"psutil: Terminating PID {pid} for user '{user}'...")
545
+ proc.terminate()
546
+ try:
547
+ proc.wait(timeout=3) # Wait a bit for termination
548
+ self.log_info(f"psutil: PID {pid} terminated successfully.")
549
+ stopped_count_psutil += 1
550
+ except psutil.TimeoutExpired:
551
+ self.log_warning(f"psutil: PID {pid} did not terminate within timeout. Will try taskkill.")
552
+ # No error append yet, let taskkill try
553
+ except (psutil.NoSuchProcess, psutil.AccessDenied):
554
+ self.log_info(f"psutil: PID {pid} already gone or access denied after terminate.")
555
+ stopped_count_psutil += 1 # Count it as stopped
556
+ except (psutil.NoSuchProcess, psutil.AccessDenied) as term_err:
557
+ self.log_warning(f"psutil: Error terminating PID {proc.pid if 'proc' in locals() and proc else 'unknown'}: {term_err}. It might be gone already.")
558
+ # If it's gone, count it?
559
+ stopped_count_psutil += 1 # Assume it's gone if NoSuchProcess
560
+ except Exception as term_ex:
561
+ self.log_error(f"psutil: Unexpected error terminating PID {proc.pid if 'proc' in locals() and proc else 'unknown'}: {term_ex}")
562
+ errors.append(f"failed_psutil_terminate_{user}")
563
+
564
+ # --- Attempt 2: taskkill fallback (for users where psutil didn't find/stop) ---
565
+ # Only run taskkill if psutil didn't stop anything for a specific user OR if target was specific user
566
+
567
+ if not self.target_executable_path:
568
+ errors.append("skipped_taskkill_no_exe_path")
569
+ else:
570
+ executable_name = os.path.basename(self.target_executable_path.strip('"'))
571
+ for user in target_users_lower:
572
+ run_taskkill = False
573
+ if user not in procs_to_kill_by_user or not procs_to_kill_by_user[user]:
574
+ # psutil didn't find anything for this user initially
575
+ run_taskkill = True
576
+ self.log_info(f"taskkill: psutil found no processes for '{user}', attempting taskkill as fallback.")
577
+ elif any(p.is_running() for p in procs_to_kill_by_user.get(user, []) if p): # Check if any psutil targets still running
578
+ run_taskkill = True
579
+ self.log_info(f"taskkill: Some processes for '{user}' may remain after psutil, attempting taskkill cleanup.")
580
+
581
+ if run_taskkill:
582
+ # Construct taskkill command
583
+ # Username format might need adjustment (e.g., DOMAIN\user). Try simple first.
584
+ taskkill_command = [
585
+ "taskkill", "/F", # Force
586
+ "/IM", executable_name, # Image name
587
+ "/FI", f"USERNAME eq {user}" # Filter by username
588
+ ]
589
+ self.log_info(f"Running taskkill command: {' '.join(taskkill_command)}")
590
+ try:
591
+ result = subprocess.run(taskkill_command, capture_output=True, text=True, check=False)
592
+ if result.returncode == 0:
593
+ self.log_info(f"taskkill successful for user '{user}' (Exit Code 0).")
594
+ # Can't easily count how many were killed here, assume success if exit 0
595
+ stopped_count_taskkill += 1 # Indicate taskkill ran successfully for user
596
+ elif result.returncode == 128: # Code 128: No tasks found matching criteria
597
+ self.log_info(f"taskkill: No matching processes found for user '{user}'.")
598
+ else:
599
+ self.log_error(f"taskkill failed for user '{user}' (Exit Code {result.returncode}).")
600
+ self.log_error(f" taskkill STDOUT: {result.stdout.strip()}")
601
+ self.log_error(f" taskkill STDERR: {result.stderr.strip()}")
602
+ errors.append(f"failed_taskkill_{user}")
603
+ except FileNotFoundError:
604
+ self.log_error("taskkill command not found.")
605
+ errors.append("failed_taskkill_not_found")
606
+ break # Stop trying taskkill if command is missing
607
+ except Exception as tk_ex:
608
+ self.log_error(f"Exception running taskkill for '{user}': {tk_ex}")
609
+ errors.append(f"failed_taskkill_exception_{user}")
610
+
611
+ # --- Consolidate status ---
612
+ final_status = "failed" # Default to failed if errors occurred
613
+ if stopped_count_psutil > 0 or stopped_count_taskkill > 0:
614
+ final_status = "success" if not errors else "partial_success"
615
+ elif not errors:
616
+ final_status = "success_no_processes_found"
617
+
618
+ details = f"psutil_stopped={stopped_count_psutil}, taskkill_users_attempted={stopped_count_taskkill}, errors={len(errors)}"
619
+ self.log_info(f"Finished stopping OOTB. Status: {final_status}. Details: {details}")
620
+ return f"{final_status}::{details}" # Return status and details
621
+
622
+ def handle_start(self, target_user="all_active"):
623
+ """Handles external start command request (finds users, calls internal trigger)."""
624
+ self.log_info(f"External start requested for target '{target_user}'...")
625
+ # This function now primarily identifies target users and calls the internal trigger method.
626
+ # The status returned here reflects the process of identifying and triggering,
627
+ # not necessarily the final success/failure of the actual start (which happens async).
628
+
629
+ active_sessions = {} # user_lower: session_id
630
+ all_system_users = set() # user_lower
631
+ try:
632
+ # Use psutil for system user list, WTS for active sessions/IDs
633
+ for user_session in psutil.users():
634
+ username_lower = user_session.name.split('\\')[-1].lower()
635
+ all_system_users.add(username_lower)
636
+
637
+ sessions = win32ts.WTSEnumerateSessions(win32ts.WTS_CURRENT_SERVER_HANDLE)
638
+ for session in sessions:
639
+ if session['State'] == win32ts.WTSActive:
640
+ try:
641
+ user = win32ts.WTSQuerySessionInformation(win32ts.WTS_CURRENT_SERVER_HANDLE, session['SessionId'], win32ts.WTSUserName)
642
+ if user:
643
+ active_sessions[user.lower()] = session['SessionId']
644
+ except Exception as query_err:
645
+ self.log_error(f"Could not query session {session['SessionId']}: {query_err}")
646
+ except Exception as user_enum_err:
647
+ self.log_error(f"Error enumerating users/sessions: {user_enum_err}", exc_info=True)
648
+ return "failed_user_enumeration"
649
+
650
+ target_users_normalized = set()
651
+ if target_user == "all_active":
652
+ target_users_normalized = set(active_sessions.keys())
653
+ self.log_info(f"Targeting all active users for start: {target_users_normalized}")
654
+ else:
655
+ normalized_target = target_user.lower()
656
+ target_users_normalized.add(normalized_target)
657
+ self.log_info(f"Targeting specific user for start: {normalized_target}")
658
+
659
+ if not target_users_normalized:
660
+ self.log_info("No target users identified for start.")
661
+ return "failed_no_target_users"
662
+
663
+ trigger_results = {}
664
+ for user in target_users_normalized:
665
+ self.log_info(f"Calling internal start trigger for user: {user}")
666
+ # Call the core logic directly (this is now synchronous within the handler)
667
+ # Or queue it? Queuing might be better to avoid blocking the handler if many users.
668
+ # Let's stick to the queue approach from the internal endpoint:
669
+ internal_command = {
670
+ "action": "_internal_start_ootb",
671
+ "target_user": user
672
+ }
673
+ internal_command_id = f"external_{user}_{time.time():.0f}"
674
+ self.command_queue.put((internal_command_id, internal_command))
675
+ trigger_results[user] = "queued"
676
+
677
+ self.log_info(f"Finished queuing start triggers. Results: {trigger_results}")
678
+ # The status here just means we successfully queued the actions
679
+ # Actual success/failure happens in the command processor later.
680
+ # We might need a different way to report overall status if needed.
681
+ return f"success_queued::{json.dumps(trigger_results)}"
682
+
683
+
684
+ def _trigger_start_for_user(self, username):
685
+ """Core logic to start OOTB for a single user. Called internally."""
686
+ user = username.lower() # Ensure lowercase
687
+ self.log_info(f"Internal trigger: Starting OOTB check for user '{user}'...")
688
+ task_created_status = "task_unknown"
689
+ immediate_start_status = "start_not_attempted"
690
+ final_status = "failed_unknown"
691
+
692
+ try:
693
+ # 1. Ensure scheduled task exists (still useful fallback/persistence)
694
+ try:
695
+ task_created = self.create_or_update_logon_task(user)
696
+ task_created_status = "task_success" if task_created else "task_failed"
697
+ except Exception as task_err:
698
+ self.log_error(f"Internal trigger: Exception creating/updating task for {user}: {task_err}", exc_info=True)
699
+ task_created_status = "task_exception"
700
+ # Don't necessarily fail the whole operation yet
701
+
702
+ # 2. Check if user is active
703
+ active_sessions = {} # Re-check active sessions specifically for this user
704
+ session_id = None
705
+ token = None
706
+ is_active = False
707
+ try:
708
+ sessions = win32ts.WTSEnumerateSessions(win32ts.WTS_CURRENT_SERVER_HANDLE)
709
+ for session in sessions:
710
+ if session['State'] == win32ts.WTSActive:
711
+ try:
712
+ current_user = win32ts.WTSQuerySessionInformation(win32ts.WTS_CURRENT_SERVER_HANDLE, session['SessionId'], win32ts.WTSUserName)
713
+ if current_user and current_user.lower() == user:
714
+ session_id = session['SessionId']
715
+ is_active = True
716
+ self.log_info(f"Internal trigger: User '{user}' is active in session {session_id}.")
717
+ break
718
+ except Exception: pass # Ignore errors querying other sessions
719
+ except Exception as e:
720
+ self.log_error(f"Internal trigger: Error checking active sessions for {user}: {e}")
721
+ # Continue, assume inactive if check failed?
722
+
723
+ if not is_active:
724
+ self.log_info(f"Internal trigger: User '{user}' is not active. Skipping immediate start.")
725
+ immediate_start_status = "start_skipped_inactive"
726
+ final_status = task_created_status # Status depends only on task creation
727
+ return final_status # Exit early if inactive
728
+
729
+ # 3. Check if already running for this active user
730
+ is_running = False
731
+ try:
732
+ running_procs = self._get_ootb_processes(user)
733
+ if running_procs:
734
+ is_running = True
735
+ self.log_info(f"Internal trigger: OOTB already running for active user '{user}'. Skipping immediate start.")
736
+ immediate_start_status = "start_skipped_already_running"
737
+ final_status = "success_already_running" # Considered success
738
+ return final_status # Exit early if already running
739
+ except Exception as e:
740
+ self.log_error(f"Internal trigger: Error checking existing processes for {user}: {e}")
741
+ # Continue and attempt start despite error?
742
+
743
+ # 4. Attempt immediate start (User is active and not running)
744
+ immediate_start_status = "start_attempted"
745
+ self.log_info(f"Internal trigger: User '{user}' is active and not running. Attempting immediate start via CreateProcessAsUser...")
746
+ try:
747
+ token = win32ts.WTSQueryUserToken(session_id)
748
+ env = win32profile.CreateEnvironmentBlock(token, False)
749
+ startup = win32process.STARTUPINFO()
750
+ creation_flags = 0x00000010 # CREATE_NEW_CONSOLE
751
+ lpApplicationName = None
752
+ lpCommandLine = f'cmd.exe /K "{self.target_executable_path}"'
753
+ cwd = os.path.dirname(self.target_executable_path.strip('"')) if os.path.dirname(self.target_executable_path.strip('"')) != '' else None
754
+
755
+ # Log details before call
756
+ self.log_info(f"Internal trigger: Calling CreateProcessAsUser:")
757
+ self.log_info(f" lpCommandLine: {lpCommandLine}")
758
+ self.log_info(f" lpCurrentDirectory: {cwd if cwd else 'Default'}")
759
+
760
+ hProcess, hThread, dwPid, dwTid = win32process.CreateProcessAsUser(
761
+ token, lpApplicationName, lpCommandLine, None, None, False,
762
+ creation_flags, env, cwd, startup
763
+ )
764
+ self.log_info(f"Internal trigger: CreateProcessAsUser call succeeded for user '{user}' (PID: {dwPid}). Checking existence...")
765
+ win32api.CloseHandle(hProcess)
766
+ win32api.CloseHandle(hThread)
767
+
768
+ time.sleep(1)
769
+ if psutil.pid_exists(dwPid):
770
+ self.log_info(f"Internal trigger: Immediate start succeeded for user '{user}' (PID {dwPid}).")
771
+ immediate_start_status = "start_success"
772
+ final_status = "success" # Overall success
773
+ else:
774
+ self.log_error(f"Internal trigger: Immediate start failed for user '{user}': Process {dwPid} exited immediately.")
775
+ immediate_start_status = "start_failed_exited"
776
+ final_status = "failed_start_exited"
777
+
778
+ except Exception as proc_err:
779
+ self.log_error(f"Internal trigger: Exception during CreateProcessAsUser for user '{user}': {proc_err}", exc_info=True)
780
+ immediate_start_status = "start_failed_exception"
781
+ final_status = "failed_start_exception"
782
+ finally:
783
+ if token: win32api.CloseHandle(token)
784
+
785
+ # Combine results (though mostly determined by start attempt now)
786
+ # Example: final_status = f"{task_created_status}_{immediate_start_status}"
787
+ return final_status
788
+
789
+ except Exception as e:
790
+ self.log_error(f"Internal trigger: Unexpected error processing start for {username}: {e}", exc_info=True)
791
+ return "failed_trigger_exception"
792
+
793
+
794
+ def create_or_update_logon_task(self, username):
795
+ """Creates/updates task to trigger the internal signal script on session connect."""
796
+ if not self.signal_script_path:
797
+ self.log_error(f"Cannot create task for {username}: Signal script path is not set.")
798
+ return False
799
+ if not sys.executable:
800
+ self.log_error(f"Cannot create task for {username}: sys.executable is not found.")
801
+ return False
802
+
803
+ # Use the python executable that the service itself is running under
804
+ python_exe = sys.executable
805
+ if ' ' in python_exe and not python_exe.startswith('"'):
806
+ python_exe = f'"{python_exe}"'
807
+
808
+ task_name = f"OOTB_UserConnect_{username}"
809
+ # Action: Run python.exe with the signal script and username argument
810
+ action_executable = python_exe
811
+ # Ensure script path is quoted if needed
812
+ script_arg = self.signal_script_path # Should be quoted already by _find_signal_script
813
+ # Username might need quoting if it contains spaces, though unlikely
814
+ user_arg = username # Keep simple for now
815
+ action_arguments = f'{script_arg} "{user_arg}"' # Pass username as quoted arg
816
+ safe_action_executable = action_executable.replace("'", "''") # Escape for PS
817
+ safe_action_arguments = action_arguments.replace("'", "''") # Escape for PS
818
+
819
+ # Working directory for the script (likely its own directory)
820
+ try:
821
+ script_dir = os.path.dirname(self.signal_script_path.strip('"'))
822
+ if not script_dir: script_dir = "."
823
+ safe_working_directory = script_dir.replace("'", "''")
824
+ working_directory_setting = f"$action.WorkingDirectory = '{safe_working_directory}'"
825
+ except Exception as e:
826
+ self.log_error(f"Error determining working directory for signal script task: {e}. WD will not be set.")
827
+ working_directory_setting = "# Could not set WorkingDirectory"
828
+
829
+ # PowerShell command construction
830
+ ps_command = f"""
831
+ $taskName = "{task_name}"
832
+ $principal = New-ScheduledTaskPrincipal -UserId "{username}" -LogonType Interactive
833
+
834
+ # Action: Run python signal script
835
+ $action = New-ScheduledTaskAction -Execute '{safe_action_executable}' -Argument '{safe_action_arguments}'
836
+ {working_directory_setting}
837
+
838
+ # Trigger: On session connect (Event ID 21)
839
+ $logName = 'Microsoft-Windows-TerminalServices-LocalSessionManager/Operational'
840
+ $source = 'Microsoft-Windows-TerminalServices-LocalSessionManager'
841
+ $eventIDs = @(21, 25)
842
+ $trigger = New-ScheduledTaskTrigger -Event -LogName $logName -Source $source -EventId $eventIDs[0]
843
+ # Optional Delay: -Delay 'PT15S'
844
+ # $trigger.Delay = 'PT15S'
845
+
846
+ $settings = New-ScheduledTaskSettingsSet -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries -ExecutionTimeLimit (New-TimeSpan -Days 9999) -RestartCount 3 -RestartInterval (New-TimeSpan -Minutes 1)
847
+ $description = "Triggers OOTB Guard Service for user {username} upon session connect via internal signal." # Updated description
848
+
849
+ # Unregister existing task first (force)
850
+ Unregister-ScheduledTask -TaskName $taskName -Confirm:$false -ErrorAction SilentlyContinue
851
+
852
+ # Register the new task
853
+ Register-ScheduledTask -TaskName $taskName -Principal $principal -Action $action -Trigger $trigger -Settings $settings -Description $description -Force
854
+ """
855
+ self.log_info(f"Attempting to create/update task '{task_name}' for user '{username}' to run signal script.")
856
+ try:
857
+ success = self.run_powershell_command(ps_command)
858
+ if success:
859
+ self.log_info(f"Successfully ran PowerShell command to create/update task '{task_name}'.")
860
+ return True
861
+ else:
862
+ self.log_error(f"PowerShell command failed to create/update task '{task_name}'. See previous logs.")
863
+ return False
864
+ except Exception as e:
865
+ self.log_error(f"Failed to create/update scheduled task '{task_name}' for user '{username}': {e}", exc_info=True)
866
+ return False
867
+
868
+ def run_powershell_command(self, command, log_output=True):
869
+ """Executes a PowerShell command and handles output/errors. Returns True on success."""
870
+ self.log_info(f"Executing PowerShell: {command}")
871
+ try:
872
+ result = subprocess.run(
873
+ ["powershell.exe", "-NoProfile", "-NonInteractive", "-Command", command],
874
+ capture_output=True, text=True, check=True, encoding='utf-8', errors='ignore'
875
+ )
876
+ if log_output and result.stdout:
877
+ self.log_info(f"PowerShell STDOUT:\n{result.stdout.strip()}")
878
+ if log_output and result.stderr:
879
+ # Log stderr as info, as some commands write status here (like unregister task not found)
880
+ self.log_info(f"PowerShell STDERR:\n{result.stderr.strip()}")
881
+ return True
882
+ except FileNotFoundError:
883
+ self.log_error("'powershell.exe' not found. Cannot manage scheduled tasks.")
884
+ return False
885
+ except subprocess.CalledProcessError as e:
886
+ # Log error but still return False, handled by caller
887
+ self.log_error(f"PowerShell command failed (Exit Code {e.returncode}):")
888
+ self.log_error(f" Command: {e.cmd}")
889
+ if e.stdout: self.log_error(f" STDOUT: {e.stdout.strip()}")
890
+ if e.stderr: self.log_error(f" STDERR: {e.stderr.strip()}")
891
+ return False
892
+ except Exception as e:
893
+ self.log_error(f"Unexpected error running PowerShell: {e}", exc_info=True)
894
+ return False
895
+
896
+ def remove_logon_task(self, username):
897
+ """Removes the logon scheduled task for a user."""
898
+ task_name = f"{self._task_name_prefix}{username}"
899
+ safe_task_name = task_name.replace("'", "''")
900
+ command = f"Unregister-ScheduledTask -TaskName '{safe_task_name}' -Confirm:$false -ErrorAction SilentlyContinue"
901
+ self.run_powershell_command(command, log_output=False)
902
+ self.log_info(f"Attempted removal of scheduled task '{task_name}' for user '{username}'.")
903
+ return True
904
+
905
+ def _find_signal_script(self):
906
+ """Finds the signal_connection.py script relative to this service file."""
907
+ try:
908
+ base_dir = os.path.dirname(os.path.abspath(__file__))
909
+ script_path = os.path.join(base_dir, "signal_connection.py")
910
+ if os.path.exists(script_path):
911
+ self.log_info(f"Found signal script at: {script_path}")
912
+ # Quote if needed?
913
+ if " " in script_path and not script_path.startswith('"'):
914
+ return f'"{script_path}"'
915
+ return script_path
916
+ else:
917
+ self.log_error(f"Signal script signal_connection.py not found near {base_dir}")
918
+ return None
919
+ except Exception as e:
920
+ self.log_error(f"Error finding signal script: {e}")
921
+ return None
922
+
923
+ # --- Main Execution Block ---
924
+ if __name__ == '__main__':
925
+ if len(sys.argv) > 1 and sys.argv[1] == 'debug':
926
+ self.log_info("Starting service in debug mode...")
927
+ print(f"Running Flask server via Waitress on {_LISTEN_HOST}:{_LISTEN_PORT} for debugging...")
928
+ print("Service logic (command processing) will NOT run in this mode.")
929
+ print("Use this primarily to test the '/command' endpoint receiving POSTs.")
930
+ print("Press Ctrl+C to stop.")
931
+ try:
932
+ serve(flask_app, host=_LISTEN_HOST, port=_LISTEN_PORT, threads=1)
933
+ except KeyboardInterrupt:
934
+ print("\nDebug server stopped.")
935
+
936
+ elif len(sys.argv) == 1:
937
+ try:
938
+ servicemanager.Initialize()
939
+ servicemanager.PrepareToHostSingle(GuardService)
940
+ servicemanager.StartServiceCtrlDispatcher()
941
+ except win32service.error as details:
942
+ import winerror
943
+ if details.winerror == winerror.ERROR_FAILED_SERVICE_CONTROLLER_CONNECT:
944
+ print(f"Error: Not started by SCM.")
945
+ print(f"Use 'python {os.path.basename(__file__)} install|start|stop|remove|debug'")
946
+ else:
947
+ print(f"Error preparing service: {details}")
948
+ except Exception as e:
949
+ print(f"Unexpected error initializing service: {e}")
950
+ else:
951
951
  win32serviceutil.HandleCommandLine(GuardService)