computer-use-ootb-internal 0.0.136__py3-none-any.whl → 0.0.137__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,836 +1,836 @@
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
- # --- Helper Functions --- Only logging helpers needed adjustments
90
- # Move these inside the class later
91
- # def get_python_executable(): ...
92
- # def get_pip_executable(): ...
93
-
94
- # Define loggers at module level for use before instance exists?
95
- # Or handle carefully within instance methods.
96
-
97
- # --- PowerShell Task Scheduler Helpers --- (These will become methods) ---
98
-
99
- # _TASK_NAME_PREFIX = "OOTB_UserLogon_" # Move to class
100
-
101
- # def run_powershell_command(command, log_output=True): ...
102
- # def create_or_update_logon_task(username, task_command, python_executable): ...
103
- # def remove_logon_task(username): ...
104
-
105
- # --- End PowerShell Task Scheduler Helpers ---
106
-
107
- TARGET_EXECUTABLE_NAME = "computer-use-ootb-internal.exe"
108
-
109
- class GuardService(win32serviceutil.ServiceFramework):
110
- _svc_name_ = _SERVICE_NAME
111
- _svc_display_name_ = _SERVICE_DISPLAY_NAME
112
- _svc_description_ = _SERVICE_DESCRIPTION
113
- _task_name_prefix = "OOTB_UserLogon_" # Class attribute for task prefix
114
-
115
- # --- Instance Logging Methods ---
116
- def log_info(self, msg):
117
- thread_name = threading.current_thread().name
118
- full_msg = f"[{thread_name}] {msg}"
119
- logging.info(full_msg)
120
- try:
121
- if threading.current_thread().name in ["MainThread", "CommandProcessor"]:
122
- servicemanager.LogInfoMsg(str(full_msg))
123
- except Exception as e:
124
- # Log only to file if event log fails
125
- logging.warning(f"(Instance) Could not write info to Windows Event Log: {e}")
126
-
127
- def log_error(self, msg, exc_info=False):
128
- thread_name = threading.current_thread().name
129
- full_msg = f"[{thread_name}] {msg}"
130
- logging.error(full_msg, exc_info=exc_info)
131
- try:
132
- if threading.current_thread().name in ["MainThread", "CommandProcessor"]:
133
- servicemanager.LogErrorMsg(str(full_msg))
134
- except Exception as e:
135
- logging.warning(f"(Instance) Could not write error to Windows Event Log: {e}")
136
- # --- End Instance Logging ---
137
-
138
- # --- Instance Helper Methods (Moved from module level) ---
139
- def _find_target_executable(self):
140
- """Finds the target executable (e.g., computer-use-ootb-internal.exe) in the Scripts directory."""
141
- try:
142
- # sys.executable should be python.exe or pythonservice.exe in the env root/Scripts
143
- env_dir = os.path.dirname(sys.executable)
144
- # If sys.executable is in Scripts, go up one level
145
- if os.path.basename(env_dir.lower()) == 'scripts':
146
- env_dir = os.path.dirname(env_dir)
147
-
148
- scripts_dir = os.path.join(env_dir, 'Scripts')
149
- target_exe_path = os.path.join(scripts_dir, TARGET_EXECUTABLE_NAME)
150
-
151
- self.log_info(f"_find_target_executable: Checking for executable at: {target_exe_path}")
152
-
153
- if os.path.exists(target_exe_path):
154
- self.log_info(f"_find_target_executable: Found executable: {target_exe_path}")
155
- # Quote if necessary for command line usage
156
- if " " in target_exe_path and not target_exe_path.startswith('"'):
157
- return f'"{target_exe_path}"'
158
- return target_exe_path
159
- else:
160
- self.log_error(f"_find_target_executable: Target executable not found at {target_exe_path}")
161
- # Fallback: Check env root directly (less common for scripts)
162
- target_exe_path_root = os.path.join(env_dir, TARGET_EXECUTABLE_NAME)
163
- self.log_info(f"_find_target_executable: Checking fallback location: {target_exe_path_root}")
164
- if os.path.exists(target_exe_path_root):
165
- self.log_info(f"_find_target_executable: Found executable at fallback location: {target_exe_path_root}")
166
- if " " in target_exe_path_root and not target_exe_path_root.startswith('"'):
167
- return f'"{target_exe_path_root}"'
168
- return target_exe_path_root
169
- else:
170
- self.log_error(f"_find_target_executable: Target executable also not found at {target_exe_path_root}")
171
- return None
172
-
173
- except Exception as e:
174
- self.log_error(f"Error finding target executable: {e}")
175
- return None
176
-
177
- def __init__(self, args):
178
- global _service_instance
179
- win32serviceutil.ServiceFramework.__init__(self, args)
180
- self.hWaitStop = win32event.CreateEvent(None, 0, 0, None)
181
- self.is_running = True
182
- self.server_thread = None
183
- self.command_queue = queue.Queue()
184
- self.command_processor_thread = None
185
- self.session = requests.Session()
186
-
187
- self.target_executable_path = self._find_target_executable()
188
- if not self.target_executable_path:
189
- # Log error and potentially stop service if critical executable is missing
190
- self.log_error(f"CRITICAL: Could not find {TARGET_EXECUTABLE_NAME}. Service cannot function.")
191
- # Consider stopping the service here if needed, or handle appropriately
192
- else:
193
- self.log_info(f"Using target executable: {self.target_executable_path}")
194
-
195
- _service_instance = self
196
- self.log_info(f"Service initialized. Target executable set to: {self.target_executable_path}")
197
-
198
- def SvcStop(self):
199
- self.log_info(f"Service stop requested.")
200
- self.ReportServiceStatus(win32service.SERVICE_STOP_PENDING)
201
- self.is_running = False
202
- # Signal the command processor thread to stop
203
- self.command_queue.put(None) # Sentinel value
204
- # Signal the main wait loop
205
- win32event.SetEvent(self.hWaitStop)
206
- # Stopping waitress gracefully from another thread is non-trivial.
207
- # We rely on the SCM timeout / process termination for now.
208
- self.log_info(f"{_SERVICE_NAME} SvcStop: Stop signaled. Server thread will be terminated by SCM.")
209
-
210
-
211
- def SvcDoRun(self):
212
- servicemanager.LogMsg(servicemanager.EVENTLOG_INFORMATION_TYPE,
213
- servicemanager.PYS_SERVICE_STARTED,
214
- (self._svc_name_, ''))
215
- try:
216
- self.log_info(f"{_SERVICE_NAME} starting.")
217
- # Start the command processor thread
218
- self.command_processor_thread = threading.Thread(
219
- target=self.process_commands, name="CommandProcessor", daemon=True)
220
- self.command_processor_thread.start()
221
- self.log_info("Command processor thread started.")
222
-
223
- # Start the Flask server (via Waitress) in a separate thread
224
- self.server_thread = threading.Thread(
225
- target=self.run_server, name="WebServerThread", daemon=True)
226
- self.server_thread.start()
227
- self.log_info(f"Web server thread started, listening on {_LISTEN_HOST}:{_LISTEN_PORT}.")
228
-
229
- self.log_info(f"{_SERVICE_NAME} running. Waiting for stop signal.")
230
- # Keep the main service thread alive waiting for stop signal
231
- win32event.WaitForSingleObject(self.hWaitStop, win32event.INFINITE)
232
- self.log_info(f"{_SERVICE_NAME} received stop signal in main thread.")
233
-
234
- except Exception as e:
235
- self.log_error(f"Fatal error in SvcDoRun: {e}", exc_info=True)
236
- self.SvcStop() # Signal stop if possible
237
- finally:
238
- self.log_info(f"{_SERVICE_NAME} SvcDoRun finished.")
239
-
240
-
241
- def run_server(self):
242
- """Runs the Flask app using Waitress."""
243
- self.log_info(f"Waitress server starting on {_LISTEN_HOST}:{_LISTEN_PORT}")
244
- try:
245
- serve(flask_app, host=_LISTEN_HOST, port=_LISTEN_PORT, threads=4)
246
- self.log_info("Waitress server has stopped.") # Should only happen on shutdown
247
- except Exception as e:
248
- self.log_error(f"Web server thread encountered an error: {e}", exc_info=True)
249
- # Consider signaling the main thread to stop if the web server fails critically
250
- # For now, just log the error.
251
-
252
-
253
- def process_commands(self):
254
- """Worker thread to process commands from the queue."""
255
- self.log_info("Command processor thread starting.")
256
- while self.is_running:
257
- try:
258
- item = self.command_queue.get(block=True, timeout=1) # Add timeout to check is_running periodically
259
- if item is None:
260
- self.log_info("Command processor received stop signal.")
261
- break # Exit loop
262
-
263
- command_id, command = item
264
- action = command.get("action")
265
- target = command.get("target_user", "all_active")
266
- status = "failed_unknown" # Default
267
-
268
- self.log_info(f"Dequeued Command ID {command_id}: action='{action}', target='{target}'")
269
-
270
- try:
271
- if action == "update":
272
- status = self.handle_update()
273
- elif action == "stop_ootb":
274
- status = self.handle_stop(target)
275
- elif action == "start_ootb":
276
- status = self.handle_start(target)
277
- else:
278
- self.log_error(f"Unknown action in queue: {action}")
279
- status = "failed_unknown_action"
280
- except Exception as handler_ex:
281
- self.log_error(f"Exception processing Command ID {command_id} ({action}): {handler_ex}", exc_info=True)
282
- status = "failed_exception"
283
- finally:
284
- self.report_command_status(command_id, status)
285
- self.command_queue.task_done()
286
-
287
- except queue.Empty:
288
- # Timeout occurred, just loop again and check self.is_running
289
- continue
290
- except Exception as e:
291
- self.log_error(f"Error in command processing loop: {e}", exc_info=True)
292
- if self.is_running:
293
- time.sleep(5)
294
-
295
- self.log_info("Command processor thread finished.")
296
-
297
-
298
- def report_command_status(self, command_id, status, details=""):
299
- """Sends command status back to the server."""
300
- if not _SERVER_STATUS_REPORT_URL:
301
- self.log_error("No server status report URL configured. Skipping report.")
302
- return
303
-
304
- payload = {
305
- "command_id": command_id,
306
- "status": status,
307
- "details": details,
308
- "machine_id": os.getenv('COMPUTERNAME', 'unknown_guard')
309
- }
310
- self.log_info(f"Reporting status for command {command_id}: {status}")
311
- try:
312
- response = self.session.post(_SERVER_STATUS_REPORT_URL, json=payload, timeout=15)
313
- response.raise_for_status()
314
- self.log_info(f"Status report for command {command_id} accepted by server.")
315
- except requests.exceptions.RequestException as e:
316
- self.log_error(f"Failed to report status for command {command_id}: {e}")
317
- except Exception as e:
318
- self.log_error(f"Unexpected error reporting status for command {command_id}: {e}", exc_info=True)
319
-
320
- # --- Command Handlers --- Now call self. for helpers
321
-
322
- def handle_update(self):
323
- self.log_info("Executing OOTB update...")
324
- if not self.target_executable_path:
325
- self.log_error("Cannot update: Target executable not found.")
326
- return "failed_executable_not_found"
327
-
328
- update_command = f"{self.target_executable_path} update"
329
- self.log_info(f"Running update command: {update_command}")
330
- try:
331
- result = subprocess.run(update_command, shell=True, capture_output=True, text=True, check=True, timeout=300, encoding='utf-8')
332
- self.log_info(f"Update successful: \nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}")
333
- return "success"
334
- except subprocess.CalledProcessError as e:
335
- self.log_error(f"Update failed (Exit Code {e.returncode}):\nSTDOUT:\n{e.stdout}\nSTDERR:\n{e.stderr}")
336
- return f"failed_exit_{e.returncode}"
337
- except subprocess.TimeoutExpired:
338
- self.log_error(f"Update command timed out.")
339
- return "failed_timeout"
340
- except Exception as e:
341
- self.log_error(f"Unexpected error during update: {e}", exc_info=True)
342
- return "failed_exception"
343
-
344
- def _get_ootb_processes(self, target_user="all_active"):
345
- ootb_procs = []
346
- target_pid_list = []
347
- try:
348
- target_users = set()
349
- if target_user == "all_active":
350
- for user_session in psutil.users():
351
- username = user_session.name.split('\\')[-1]
352
- target_users.add(username.lower())
353
- else:
354
- target_users.add(target_user.lower())
355
- self.log_info(f"Searching for OOTB processes for users: {target_users}")
356
-
357
- # Use the potentially corrected python.exe path for matching
358
- python_exe_path_for_check = self.target_executable_path.strip('"')
359
- self.log_info(f"_get_ootb_processes: Checking against python path: {python_exe_path_for_check}")
360
-
361
- for proc in psutil.process_iter(['pid', 'name', 'username', 'cmdline', 'exe']):
362
- try:
363
- pinfo = proc.info
364
- proc_username = pinfo['username']
365
- if proc_username:
366
- proc_username = proc_username.split('\\')[-1].lower()
367
-
368
- if proc_username in target_users:
369
- cmdline = ' '.join(pinfo['cmdline']) if pinfo['cmdline'] else ''
370
- # Check if the process executable matches our corrected python path AND module is in cmdline
371
- if pinfo['exe'] and pinfo['exe'] == python_exe_path_for_check and _OOTB_MODULE in cmdline:
372
- self.log_info(f"Found matching OOTB process: PID={pinfo['pid']}, User={pinfo['username']}, Cmd={cmdline}")
373
- ootb_procs.append(proc)
374
- target_pid_list.append(pinfo['pid'])
375
- except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
376
- continue
377
- self.log_info(f"Found {len(ootb_procs)} OOTB process(es) matching criteria: {target_pid_list}")
378
- except Exception as e:
379
- self.log_error(f"Error enumerating processes: {e}", exc_info=True)
380
- return ootb_procs
381
-
382
- def handle_stop(self, target_user="all_active"):
383
- self.log_info(f"Executing stop OOTB for target '{target_user}'...")
384
- stop_results = {} # Track results per user {username: (task_status, immediate_status)}
385
- failed_users = set()
386
-
387
- try:
388
- # --- Get target users and active sessions ---
389
- active_sessions = {} # user_lower: session_id
390
- # No need for all_system_users for stop, we only care about active or the specific target
391
- try:
392
- sessions = win32ts.WTSEnumerateSessions(win32ts.WTS_CURRENT_SERVER_HANDLE)
393
- for session in sessions:
394
- if session['State'] == win32ts.WTSActive:
395
- try:
396
- user = win32ts.WTSQuerySessionInformation(win32ts.WTS_CURRENT_SERVER_HANDLE, session['SessionId'], win32ts.WTSUserName)
397
- if user:
398
- active_sessions[user.lower()] = session['SessionId']
399
- except Exception as query_err:
400
- self.log_error(f"Could not query session {session['SessionId']} during stop: {query_err}")
401
- except Exception as user_enum_err:
402
- self.log_error(f"Error enumerating users/sessions during stop: {user_enum_err}", exc_info=True)
403
- return "failed_user_enumeration"
404
-
405
- self.log_info(f"Stop target: '{target_user}'. Active sessions: {active_sessions}")
406
-
407
- target_users_normalized = set()
408
- if target_user == "all_active":
409
- # Target only currently active users for stop all
410
- target_users_normalized = set(active_sessions.keys())
411
- self.log_info(f"Stop targeting all active users: {target_users_normalized}")
412
- else:
413
- # Target the specific user, regardless of active status (for task removal)
414
- normalized_target = target_user.lower()
415
- target_users_normalized.add(normalized_target)
416
- self.log_info(f"Stop targeting specific user: {normalized_target}")
417
-
418
- if not target_users_normalized:
419
- self.log_info("No target users identified for stop.")
420
- return "failed_no_target_users" # Or success if none were targeted?
421
-
422
- # --- Process each target user ---
423
- for user in target_users_normalized:
424
- task_removed_status = "task_unknown"
425
- immediate_stop_status = "stop_not_attempted"
426
- stopped_count = 0
427
-
428
- self.log_info(f"Processing stop for user '{user}'...")
429
-
430
- # 1. Always try to remove the scheduled task
431
- try:
432
- # remove_logon_task always returns True for now, just logs attempt
433
- self.remove_logon_task(user)
434
- task_removed_status = "task_removed_attempted"
435
- except Exception as task_err:
436
- self.log_error(f"Exception removing scheduled task for {user}: {task_err}", exc_info=True)
437
- task_removed_status = "task_exception"
438
- failed_users.add(user)
439
- # Continue to try and stop process if active
440
-
441
- # 2. If user is active, try to terminate process
442
- is_active = user in active_sessions
443
-
444
- if is_active:
445
- immediate_stop_status = "stop_attempted"
446
- self.log_info(f"User '{user}' is active. Attempting to terminate OOTB process(es)...")
447
- # Pass the specific username to _get_ootb_processes
448
- procs_to_stop = self._get_ootb_processes(user)
449
-
450
- if not procs_to_stop:
451
- self.log_info(f"No running OOTB processes found for active user '{user}'.")
452
- immediate_stop_status = "stop_skipped_not_running"
453
- else:
454
- self.log_info(f"Found {len(procs_to_stop)} process(es) for user '{user}' to stop.")
455
- for proc in procs_to_stop:
456
- try:
457
- pid = proc.pid # Get pid before potential termination
458
- username = proc.info.get('username', 'unknown_user')
459
- self.log_info(f"Terminating process PID={pid}, User={username}")
460
- proc.terminate()
461
- try:
462
- proc.wait(timeout=3)
463
- self.log_info(f"Process PID={pid} terminated successfully.")
464
- stopped_count += 1
465
- except psutil.TimeoutExpired:
466
- self.log_error(f"Process PID={pid} did not terminate gracefully, killing.")
467
- proc.kill()
468
- stopped_count += 1
469
- except psutil.NoSuchProcess:
470
- self.log_info(f"Process PID={pid} already terminated.")
471
- # Don't increment stopped_count here as we didn't stop it now
472
- except psutil.AccessDenied:
473
- self.log_error(f"Access denied trying to terminate process PID={pid}.")
474
- failed_users.add(user) # Mark user as failed if stop fails
475
- except Exception as e:
476
- self.log_error(f"Error stopping process PID={pid}: {e}", exc_info=True)
477
- failed_users.add(user) # Mark user as failed
478
-
479
- # Determine status based on how many were found vs stopped
480
- if user in failed_users:
481
- immediate_stop_status = f"stop_errors_terminated_{stopped_count}_of_{len(procs_to_stop)}"
482
- elif stopped_count == len(procs_to_stop):
483
- immediate_stop_status = f"stop_success_terminated_{stopped_count}"
484
- else: # Should ideally not happen if NoSuchProcess doesn't count
485
- immediate_stop_status = f"stop_partial_terminated_{stopped_count}_of_{len(procs_to_stop)}"
486
-
487
- else: # User not active
488
- self.log_info(f"User '{user}' is not active. Skipping immediate process stop (task removal attempted).")
489
- immediate_stop_status = "stop_skipped_inactive"
490
-
491
- # Record final results for this user
492
- stop_results[user] = (task_removed_status, immediate_stop_status)
493
-
494
-
495
- # --- Consolidate status ---
496
- total_processed = len(target_users_normalized)
497
- final_status = "partial_success" if failed_users else "success"
498
- if not stop_results: final_status = "no_targets_processed"
499
- if len(failed_users) == total_processed and total_processed > 0 : final_status = "failed"
500
-
501
- self.log_info(f"Finished stopping OOTB. Overall Status: {final_status}. Results: {stop_results}")
502
- try:
503
- details = json.dumps(stop_results)
504
- except Exception:
505
- details = str(stop_results) # Fallback
506
- return f"{final_status}::{details}" # Use :: as separator
507
-
508
- except Exception as e:
509
- self.log_error(f"Error during combined stop OOTB process: {e}", exc_info=True)
510
- return "failed_exception"
511
-
512
-
513
- def handle_start(self, target_user="all_active"):
514
- self.log_info(f"Executing start OOTB for target '{target_user}'...")
515
- start_results = {} # Track results per user {username: (task_status, immediate_status)}
516
- failed_users = set()
517
-
518
- try:
519
- # --- Get target users and active sessions ---
520
- active_sessions = {} # user_lower: session_id
521
- all_system_users = set() # user_lower
522
- try:
523
- # Use psutil for system user list, WTS for active sessions/IDs
524
- for user_session in psutil.users():
525
- username_lower = user_session.name.split('\\')[-1].lower()
526
- all_system_users.add(username_lower)
527
-
528
- sessions = win32ts.WTSEnumerateSessions(win32ts.WTS_CURRENT_SERVER_HANDLE)
529
- for session in sessions:
530
- if session['State'] == win32ts.WTSActive:
531
- try:
532
- user = win32ts.WTSQuerySessionInformation(win32ts.WTS_CURRENT_SERVER_HANDLE, session['SessionId'], win32ts.WTSUserName)
533
- if user:
534
- active_sessions[user.lower()] = session['SessionId']
535
- except Exception as query_err:
536
- self.log_error(f"Could not query session {session['SessionId']}: {query_err}")
537
- except Exception as user_enum_err:
538
- self.log_error(f"Error enumerating users/sessions: {user_enum_err}", exc_info=True)
539
- return "failed_user_enumeration"
540
-
541
- self.log_info(f"Found active user sessions: {active_sessions}")
542
-
543
- target_users_normalized = set()
544
- if target_user == "all_active":
545
- # If targeting all_active, only target those CURRENTLY active
546
- target_users_normalized = set(active_sessions.keys())
547
- self.log_info(f"Targeting all active users: {target_users_normalized}")
548
- else:
549
- normalized_target = target_user.lower()
550
- # Check if the target user actually exists on the system, even if inactive
551
- # This check might be complex/unreliable. Rely on task scheduler potentially failing?
552
- # Let's assume admin provides a valid username for specific targeting.
553
- # if normalized_target in all_system_users: # Removing this check, assume valid user input
554
- target_users_normalized.add(normalized_target)
555
- self.log_info(f"Targeting specific user: {normalized_target}")
556
- # else:
557
- # log_error(f"Target user '{target_user}' does not appear to exist on this system based on psutil.")
558
- # return "failed_user_does_not_exist"
559
-
560
- if not target_users_normalized:
561
- self.log_info("No target users identified (or none active for 'all_active').")
562
- # If target was specific user but they weren't found active, still try task?
563
- # Let's proceed to task creation anyway for specific user case.
564
- if target_user != "all_active": target_users_normalized.add(target_user.lower())
565
- if not target_users_normalized:
566
- return "failed_no_target_users"
567
-
568
- # --- Check existing processes ---
569
- # This check is only relevant for immediate start attempt
570
- running_procs_by_user = {} # user_lower: count
571
- try:
572
- current_running = self._get_ootb_processes("all_active") # Check all
573
- for proc in current_running:
574
- try:
575
- proc_username = proc.info.get('username')
576
- if proc_username:
577
- user_lower = proc_username.split('\\')[-1].lower()
578
- running_procs_by_user[user_lower] = running_procs_by_user.get(user_lower, 0) + 1
579
- except Exception: pass
580
- except Exception as e:
581
- self.log_error(f"Error checking existing processes: {e}")
582
- self.log_info(f"Users currently running OOTB: {running_procs_by_user}")
583
-
584
- # --- Process each target user ---
585
- for user in target_users_normalized:
586
- task_created_status = "task_unknown"
587
- immediate_start_status = "start_not_attempted"
588
- token = None # Ensure token is reset/defined
589
-
590
- self.log_info(f"Processing start for user '{user}'...")
591
-
592
- # 1. Always try to create/update the scheduled task
593
- try:
594
- task_created = self.create_or_update_logon_task(user)
595
- task_created_status = "task_success" if task_created else "task_failed"
596
- except Exception as task_err:
597
- self.log_error(f"Exception creating/updating scheduled task for {user}: {task_err}", exc_info=True)
598
- task_created_status = "task_exception"
599
- failed_users.add(user)
600
- # Continue to potentially try immediate start IF user is active?
601
- # Or maybe skip if task creation failed badly?
602
- # Let's skip immediate start if task creation had exception.
603
- start_results[user] = (task_created_status, immediate_start_status)
604
- continue
605
-
606
- # 2. If user is active AND not already running, try immediate start
607
- is_active = user in active_sessions
608
- is_running = running_procs_by_user.get(user, 0) > 0
609
-
610
- if is_active:
611
- if not is_running:
612
- immediate_start_status = "start_attempted"
613
- self.log_info(f"User '{user}' is active and not running OOTB. Attempting immediate start...")
614
- try:
615
- session_id = active_sessions[user]
616
- token = win32ts.WTSQueryUserToken(session_id)
617
- env = win32profile.CreateEnvironmentBlock(token, False) # <-- Restore creating env block
618
- startup = win32process.STARTUPINFO()
619
- creation_flags = 0x00000010 # CREATE_NEW_CONSOLE
620
-
621
- # --- Launch via cmd /K ---
622
- lpApplicationName = None
623
- lpCommandLine = f'cmd.exe /K "{self.target_executable_path}"'
624
- cwd = os.path.dirname(self.target_executable_path.strip('"')) if os.path.dirname(self.target_executable_path.strip('"')) != '' else None
625
- # --- End Launch via cmd /K ---
626
-
627
- self.log_info(f"Calling CreateProcessAsUser to launch via cmd /K (Restored Environment Block):") # Updated log note
628
- self.log_info(f" lpApplicationName: {lpApplicationName}")
629
- self.log_info(f" lpCommandLine: {lpCommandLine}")
630
- self.log_info(f" lpCurrentDirectory: {cwd if cwd else 'Default'}")
631
- self.log_info(f" dwCreationFlags: {creation_flags} (CREATE_NEW_CONSOLE)")
632
-
633
- # CreateProcessAsUser call with restored env
634
- hProcess, hThread, dwPid, dwTid = win32process.CreateProcessAsUser(
635
- token, # User token
636
- lpApplicationName, # Application name (None)
637
- lpCommandLine, # Command line (cmd.exe /K "...")
638
- None, # Process attributes
639
- None, # Thread attributes
640
- False, # Inherit handles
641
- creation_flags, # Creation flags (CREATE_NEW_CONSOLE)
642
- env, # Environment block (Restored)
643
- cwd, # Current directory for cmd.exe
644
- startup # Startup info
645
- )
646
-
647
- self.log_info(f"CreateProcessAsUser call succeeded for user '{user}' (PID: {dwPid}). Checking existence...")
648
- win32api.CloseHandle(hProcess)
649
- win32api.CloseHandle(hThread)
650
-
651
- time.sleep(1)
652
- if psutil.pid_exists(dwPid):
653
- self.log_info(f"Immediate start succeeded for user '{user}' (PID {dwPid}).")
654
- immediate_start_status = "start_success"
655
- else:
656
- self.log_error(f"Immediate start failed for user '{user}': Process {dwPid} exited immediately.")
657
- immediate_start_status = "start_failed_exited"
658
- failed_users.add(user)
659
-
660
- except Exception as proc_err:
661
- self.log_error(f"Exception during immediate start for user '{user}': {proc_err}", exc_info=True)
662
- immediate_start_status = "start_failed_exception"
663
- failed_users.add(user)
664
- finally:
665
- if token:
666
- try: win32api.CloseHandle(token)
667
- except: pass
668
- else: # User is active but already running
669
- self.log_info(f"User '{user}' is active but OOTB is already running. Skipping immediate start.")
670
- immediate_start_status = "start_skipped_already_running"
671
- else: # User is not active
672
- self.log_info(f"User '{user}' is not active. Skipping immediate start (task created/updated).")
673
- immediate_start_status = "start_skipped_inactive"
674
-
675
- # Record final results for this user
676
- start_results[user] = (task_created_status, immediate_start_status)
677
-
678
-
679
- # --- Consolidate status ---
680
- total_processed = len(target_users_normalized)
681
- final_status = "partial_success" if failed_users else "success"
682
- if not start_results: final_status = "no_targets_processed"
683
- # If all processed users failed in some way (either task or start)
684
- if len(failed_users) == total_processed and total_processed > 0: final_status = "failed"
685
- # Special case: target was specific user who wasn't found active
686
- elif total_processed == 1 and target_user != "all_active" and target_user.lower() not in active_sessions:
687
- user_key = target_user.lower()
688
- if user_key in start_results and start_results[user_key][0] == "task_success":
689
- final_status = "success_task_only_user_inactive"
690
- else:
691
- final_status = "failed_task_user_inactive"
692
-
693
- self.log_info(f"Finished starting OOTB. Overall Status: {final_status}. Results: {start_results}")
694
- # Return detailed results as a JSON string for easier parsing/logging server-side
695
- try:
696
- details = json.dumps(start_results)
697
- except Exception:
698
- details = str(start_results) # Fallback
699
- return f"{final_status}::{details}"
700
-
701
- except Exception as e:
702
- self.log_error(f"Error during combined start OOTB process: {e}", exc_info=True)
703
- return "failed_exception"
704
-
705
- def create_or_update_logon_task(self, username):
706
- """Creates/updates task to run OOTB app via cmd /K on session connect/reconnect."""
707
- if not self.target_executable_path:
708
- self.log_error(f"Cannot create task for {username}: Target executable path is not set.")
709
- return False
710
-
711
- task_name = f"OOTB_UserConnect_{username}"
712
- # Action: Revert to running the actual executable via cmd /K
713
- action_executable = 'cmd.exe'
714
- action_arguments = f'/K "{self.target_executable_path}"'
715
- safe_action_executable = action_executable.replace("'", "''")
716
- safe_action_arguments = action_arguments.replace("'", "''")
717
-
718
- # Explicitly set the working directory to the executable's location
719
- try:
720
- executable_dir = os.path.dirname(self.target_executable_path.strip('"'))
721
- if not executable_dir: executable_dir = "."
722
- safe_working_directory = executable_dir.replace("'", "''")
723
- working_directory_setting = f"$action.WorkingDirectory = '{safe_working_directory}'"
724
- except Exception as e:
725
- self.log_error(f"Error determining working directory for task: {e}. WD will not be set.")
726
- working_directory_setting = "# Could not set WorkingDirectory"
727
-
728
- # PowerShell command construction
729
- ps_command = f"""
730
- $taskName = "{task_name}"
731
- $principal = New-ScheduledTaskPrincipal -UserId "{username}" -LogonType Interactive
732
-
733
- # Action: Run the OOTB executable via cmd /K
734
- $action = New-ScheduledTaskAction -Execute '{safe_action_executable}' -Argument '{safe_action_arguments}'
735
- {working_directory_setting} # Set the working directory
736
-
737
- # Triggers: On session connect (21) AND reconnect (25)
738
- $logName = 'Microsoft-Windows-TerminalServices-LocalSessionManager/Operational'
739
- $source = 'Microsoft-Windows-TerminalServices-LocalSessionManager'
740
- # Define multiple triggers
741
- $trigger1 = New-ScheduledTaskTrigger -Event -LogName $logName -Source $source -EventId 21
742
- $trigger2 = New-ScheduledTaskTrigger -Event -LogName $logName -Source $source -EventId 25
743
- # Optional Delay - Apply to both triggers if desired?
744
- # $trigger1.Delay = 'PT15S'
745
- # $trigger2.Delay = 'PT15S'
746
-
747
- $settings = New-ScheduledTaskSettingsSet -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries -ExecutionTimeLimit (New-TimeSpan -Days 9999) -RestartCount 3 -RestartInterval (New-TimeSpan -Minutes 1)
748
- $description = "Runs OOTB Application (via cmd) for user {username} upon session connect/reconnect." # Updated description
749
-
750
- # Unregister existing task first (force) - Use the NEW task name
751
- Unregister-ScheduledTask -TaskName $taskName -Confirm:$false -ErrorAction SilentlyContinue
752
-
753
- # Register the new task with MULTIPLE triggers
754
- Register-ScheduledTask -TaskName $taskName -Principal $principal -Action $action -Trigger $trigger1, $trigger2 -Settings $settings -Description $description -Force
755
-
756
- """
757
- self.log_info(f"Attempting to create/update task '{task_name}' for user '{username}' to run OOTB on session connect/reconnect.")
758
- try:
759
- # Need to actually run the powershell command here!
760
- success = self.run_powershell_command(ps_command)
761
- if success:
762
- self.log_info(f"Successfully ran PowerShell command to create/update task '{task_name}'.")
763
- return True
764
- else:
765
- self.log_error(f"PowerShell command failed to create/update task '{task_name}'. See previous logs.")
766
- return False
767
- except Exception as e:
768
- self.log_error(f"Failed to create/update scheduled task '{task_name}' for user '{username}': {e}", exc_info=True)
769
- return False
770
-
771
- def run_powershell_command(self, command, log_output=True):
772
- """Executes a PowerShell command and handles output/errors. Returns True on success."""
773
- self.log_info(f"Executing PowerShell: {command}")
774
- try:
775
- result = subprocess.run(
776
- ["powershell.exe", "-NoProfile", "-NonInteractive", "-Command", command],
777
- capture_output=True, text=True, check=True, encoding='utf-8', errors='ignore'
778
- )
779
- if log_output and result.stdout:
780
- self.log_info(f"PowerShell STDOUT:\n{result.stdout.strip()}")
781
- if log_output and result.stderr:
782
- # Log stderr as info, as some commands write status here (like unregister task not found)
783
- self.log_info(f"PowerShell STDERR:\n{result.stderr.strip()}")
784
- return True
785
- except FileNotFoundError:
786
- self.log_error("'powershell.exe' not found. Cannot manage scheduled tasks.")
787
- return False
788
- except subprocess.CalledProcessError as e:
789
- # Log error but still return False, handled by caller
790
- self.log_error(f"PowerShell command failed (Exit Code {e.returncode}):")
791
- self.log_error(f" Command: {e.cmd}")
792
- if e.stdout: self.log_error(f" STDOUT: {e.stdout.strip()}")
793
- if e.stderr: self.log_error(f" STDERR: {e.stderr.strip()}")
794
- return False
795
- except Exception as e:
796
- self.log_error(f"Unexpected error running PowerShell: {e}", exc_info=True)
797
- return False
798
-
799
- def remove_logon_task(self, username):
800
- """Removes the logon scheduled task for a user."""
801
- task_name = f"{self._task_name_prefix}{username}"
802
- safe_task_name = task_name.replace("'", "''")
803
- command = f"Unregister-ScheduledTask -TaskName '{safe_task_name}' -Confirm:$false -ErrorAction SilentlyContinue"
804
- self.run_powershell_command(command, log_output=False)
805
- self.log_info(f"Attempted removal of scheduled task '{task_name}' for user '{username}'.")
806
- return True
807
-
808
- # --- Main Execution Block ---
809
- if __name__ == '__main__':
810
- if len(sys.argv) > 1 and sys.argv[1] == 'debug':
811
- self.log_info("Starting service in debug mode...")
812
- print(f"Running Flask server via Waitress on {_LISTEN_HOST}:{_LISTEN_PORT} for debugging...")
813
- print("Service logic (command processing) will NOT run in this mode.")
814
- print("Use this primarily to test the '/command' endpoint receiving POSTs.")
815
- print("Press Ctrl+C to stop.")
816
- try:
817
- serve(flask_app, host=_LISTEN_HOST, port=_LISTEN_PORT, threads=1)
818
- except KeyboardInterrupt:
819
- print("\nDebug server stopped.")
820
-
821
- elif len(sys.argv) == 1:
822
- try:
823
- servicemanager.Initialize()
824
- servicemanager.PrepareToHostSingle(GuardService)
825
- servicemanager.StartServiceCtrlDispatcher()
826
- except win32service.error as details:
827
- import winerror
828
- if details.winerror == winerror.ERROR_FAILED_SERVICE_CONTROLLER_CONNECT:
829
- print(f"Error: Not started by SCM.")
830
- print(f"Use 'python {os.path.basename(__file__)} install|start|stop|remove|debug'")
831
- else:
832
- print(f"Error preparing service: {details}")
833
- except Exception as e:
834
- print(f"Unexpected error initializing service: {e}")
835
- 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
+ # --- Helper Functions --- Only logging helpers needed adjustments
90
+ # Move these inside the class later
91
+ # def get_python_executable(): ...
92
+ # def get_pip_executable(): ...
93
+
94
+ # Define loggers at module level for use before instance exists?
95
+ # Or handle carefully within instance methods.
96
+
97
+ # --- PowerShell Task Scheduler Helpers --- (These will become methods) ---
98
+
99
+ # _TASK_NAME_PREFIX = "OOTB_UserLogon_" # Move to class
100
+
101
+ # def run_powershell_command(command, log_output=True): ...
102
+ # def create_or_update_logon_task(username, task_command, python_executable): ...
103
+ # def remove_logon_task(username): ...
104
+
105
+ # --- End PowerShell Task Scheduler Helpers ---
106
+
107
+ TARGET_EXECUTABLE_NAME = "computer-use-ootb-internal.exe"
108
+
109
+ class GuardService(win32serviceutil.ServiceFramework):
110
+ _svc_name_ = _SERVICE_NAME
111
+ _svc_display_name_ = _SERVICE_DISPLAY_NAME
112
+ _svc_description_ = _SERVICE_DESCRIPTION
113
+ _task_name_prefix = "OOTB_UserLogon_" # Class attribute for task prefix
114
+
115
+ # --- Instance Logging Methods ---
116
+ def log_info(self, msg):
117
+ thread_name = threading.current_thread().name
118
+ full_msg = f"[{thread_name}] {msg}"
119
+ logging.info(full_msg)
120
+ try:
121
+ if threading.current_thread().name in ["MainThread", "CommandProcessor"]:
122
+ servicemanager.LogInfoMsg(str(full_msg))
123
+ except Exception as e:
124
+ # Log only to file if event log fails
125
+ logging.warning(f"(Instance) Could not write info to Windows Event Log: {e}")
126
+
127
+ def log_error(self, msg, exc_info=False):
128
+ thread_name = threading.current_thread().name
129
+ full_msg = f"[{thread_name}] {msg}"
130
+ logging.error(full_msg, exc_info=exc_info)
131
+ try:
132
+ if threading.current_thread().name in ["MainThread", "CommandProcessor"]:
133
+ servicemanager.LogErrorMsg(str(full_msg))
134
+ except Exception as e:
135
+ logging.warning(f"(Instance) Could not write error to Windows Event Log: {e}")
136
+ # --- End Instance Logging ---
137
+
138
+ # --- Instance Helper Methods (Moved from module level) ---
139
+ def _find_target_executable(self):
140
+ """Finds the target executable (e.g., computer-use-ootb-internal.exe) in the Scripts directory."""
141
+ try:
142
+ # sys.executable should be python.exe or pythonservice.exe in the env root/Scripts
143
+ env_dir = os.path.dirname(sys.executable)
144
+ # If sys.executable is in Scripts, go up one level
145
+ if os.path.basename(env_dir.lower()) == 'scripts':
146
+ env_dir = os.path.dirname(env_dir)
147
+
148
+ scripts_dir = os.path.join(env_dir, 'Scripts')
149
+ target_exe_path = os.path.join(scripts_dir, TARGET_EXECUTABLE_NAME)
150
+
151
+ self.log_info(f"_find_target_executable: Checking for executable at: {target_exe_path}")
152
+
153
+ if os.path.exists(target_exe_path):
154
+ self.log_info(f"_find_target_executable: Found executable: {target_exe_path}")
155
+ # Quote if necessary for command line usage
156
+ if " " in target_exe_path and not target_exe_path.startswith('"'):
157
+ return f'"{target_exe_path}"'
158
+ return target_exe_path
159
+ else:
160
+ self.log_error(f"_find_target_executable: Target executable not found at {target_exe_path}")
161
+ # Fallback: Check env root directly (less common for scripts)
162
+ target_exe_path_root = os.path.join(env_dir, TARGET_EXECUTABLE_NAME)
163
+ self.log_info(f"_find_target_executable: Checking fallback location: {target_exe_path_root}")
164
+ if os.path.exists(target_exe_path_root):
165
+ self.log_info(f"_find_target_executable: Found executable at fallback location: {target_exe_path_root}")
166
+ if " " in target_exe_path_root and not target_exe_path_root.startswith('"'):
167
+ return f'"{target_exe_path_root}"'
168
+ return target_exe_path_root
169
+ else:
170
+ self.log_error(f"_find_target_executable: Target executable also not found at {target_exe_path_root}")
171
+ return None
172
+
173
+ except Exception as e:
174
+ self.log_error(f"Error finding target executable: {e}")
175
+ return None
176
+
177
+ def __init__(self, args):
178
+ global _service_instance
179
+ win32serviceutil.ServiceFramework.__init__(self, args)
180
+ self.hWaitStop = win32event.CreateEvent(None, 0, 0, None)
181
+ self.is_running = True
182
+ self.server_thread = None
183
+ self.command_queue = queue.Queue()
184
+ self.command_processor_thread = None
185
+ self.session = requests.Session()
186
+
187
+ self.target_executable_path = self._find_target_executable()
188
+ if not self.target_executable_path:
189
+ # Log error and potentially stop service if critical executable is missing
190
+ self.log_error(f"CRITICAL: Could not find {TARGET_EXECUTABLE_NAME}. Service cannot function.")
191
+ # Consider stopping the service here if needed, or handle appropriately
192
+ else:
193
+ self.log_info(f"Using target executable: {self.target_executable_path}")
194
+
195
+ _service_instance = self
196
+ self.log_info(f"Service initialized. Target executable set to: {self.target_executable_path}")
197
+
198
+ def SvcStop(self):
199
+ self.log_info(f"Service stop requested.")
200
+ self.ReportServiceStatus(win32service.SERVICE_STOP_PENDING)
201
+ self.is_running = False
202
+ # Signal the command processor thread to stop
203
+ self.command_queue.put(None) # Sentinel value
204
+ # Signal the main wait loop
205
+ win32event.SetEvent(self.hWaitStop)
206
+ # Stopping waitress gracefully from another thread is non-trivial.
207
+ # We rely on the SCM timeout / process termination for now.
208
+ self.log_info(f"{_SERVICE_NAME} SvcStop: Stop signaled. Server thread will be terminated by SCM.")
209
+
210
+
211
+ def SvcDoRun(self):
212
+ servicemanager.LogMsg(servicemanager.EVENTLOG_INFORMATION_TYPE,
213
+ servicemanager.PYS_SERVICE_STARTED,
214
+ (self._svc_name_, ''))
215
+ try:
216
+ self.log_info(f"{_SERVICE_NAME} starting.")
217
+ # Start the command processor thread
218
+ self.command_processor_thread = threading.Thread(
219
+ target=self.process_commands, name="CommandProcessor", daemon=True)
220
+ self.command_processor_thread.start()
221
+ self.log_info("Command processor thread started.")
222
+
223
+ # Start the Flask server (via Waitress) in a separate thread
224
+ self.server_thread = threading.Thread(
225
+ target=self.run_server, name="WebServerThread", daemon=True)
226
+ self.server_thread.start()
227
+ self.log_info(f"Web server thread started, listening on {_LISTEN_HOST}:{_LISTEN_PORT}.")
228
+
229
+ self.log_info(f"{_SERVICE_NAME} running. Waiting for stop signal.")
230
+ # Keep the main service thread alive waiting for stop signal
231
+ win32event.WaitForSingleObject(self.hWaitStop, win32event.INFINITE)
232
+ self.log_info(f"{_SERVICE_NAME} received stop signal in main thread.")
233
+
234
+ except Exception as e:
235
+ self.log_error(f"Fatal error in SvcDoRun: {e}", exc_info=True)
236
+ self.SvcStop() # Signal stop if possible
237
+ finally:
238
+ self.log_info(f"{_SERVICE_NAME} SvcDoRun finished.")
239
+
240
+
241
+ def run_server(self):
242
+ """Runs the Flask app using Waitress."""
243
+ self.log_info(f"Waitress server starting on {_LISTEN_HOST}:{_LISTEN_PORT}")
244
+ try:
245
+ serve(flask_app, host=_LISTEN_HOST, port=_LISTEN_PORT, threads=4)
246
+ self.log_info("Waitress server has stopped.") # Should only happen on shutdown
247
+ except Exception as e:
248
+ self.log_error(f"Web server thread encountered an error: {e}", exc_info=True)
249
+ # Consider signaling the main thread to stop if the web server fails critically
250
+ # For now, just log the error.
251
+
252
+
253
+ def process_commands(self):
254
+ """Worker thread to process commands from the queue."""
255
+ self.log_info("Command processor thread starting.")
256
+ while self.is_running:
257
+ try:
258
+ item = self.command_queue.get(block=True, timeout=1) # Add timeout to check is_running periodically
259
+ if item is None:
260
+ self.log_info("Command processor received stop signal.")
261
+ break # Exit loop
262
+
263
+ command_id, command = item
264
+ action = command.get("action")
265
+ target = command.get("target_user", "all_active")
266
+ status = "failed_unknown" # Default
267
+
268
+ self.log_info(f"Dequeued Command ID {command_id}: action='{action}', target='{target}'")
269
+
270
+ try:
271
+ if action == "update":
272
+ status = self.handle_update()
273
+ elif action == "stop_ootb":
274
+ status = self.handle_stop(target)
275
+ elif action == "start_ootb":
276
+ status = self.handle_start(target)
277
+ else:
278
+ self.log_error(f"Unknown action in queue: {action}")
279
+ status = "failed_unknown_action"
280
+ except Exception as handler_ex:
281
+ self.log_error(f"Exception processing Command ID {command_id} ({action}): {handler_ex}", exc_info=True)
282
+ status = "failed_exception"
283
+ finally:
284
+ self.report_command_status(command_id, status)
285
+ self.command_queue.task_done()
286
+
287
+ except queue.Empty:
288
+ # Timeout occurred, just loop again and check self.is_running
289
+ continue
290
+ except Exception as e:
291
+ self.log_error(f"Error in command processing loop: {e}", exc_info=True)
292
+ if self.is_running:
293
+ time.sleep(5)
294
+
295
+ self.log_info("Command processor thread finished.")
296
+
297
+
298
+ def report_command_status(self, command_id, status, details=""):
299
+ """Sends command status back to the server."""
300
+ if not _SERVER_STATUS_REPORT_URL:
301
+ self.log_error("No server status report URL configured. Skipping report.")
302
+ return
303
+
304
+ payload = {
305
+ "command_id": command_id,
306
+ "status": status,
307
+ "details": details,
308
+ "machine_id": os.getenv('COMPUTERNAME', 'unknown_guard')
309
+ }
310
+ self.log_info(f"Reporting status for command {command_id}: {status}")
311
+ try:
312
+ response = self.session.post(_SERVER_STATUS_REPORT_URL, json=payload, timeout=15)
313
+ response.raise_for_status()
314
+ self.log_info(f"Status report for command {command_id} accepted by server.")
315
+ except requests.exceptions.RequestException as e:
316
+ self.log_error(f"Failed to report status for command {command_id}: {e}")
317
+ except Exception as e:
318
+ self.log_error(f"Unexpected error reporting status for command {command_id}: {e}", exc_info=True)
319
+
320
+ # --- Command Handlers --- Now call self. for helpers
321
+
322
+ def handle_update(self):
323
+ self.log_info("Executing OOTB update...")
324
+ if not self.target_executable_path:
325
+ self.log_error("Cannot update: Target executable not found.")
326
+ return "failed_executable_not_found"
327
+
328
+ update_command = f"{self.target_executable_path} update"
329
+ self.log_info(f"Running update command: {update_command}")
330
+ try:
331
+ result = subprocess.run(update_command, shell=True, capture_output=True, text=True, check=True, timeout=300, encoding='utf-8')
332
+ self.log_info(f"Update successful: \nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}")
333
+ return "success"
334
+ except subprocess.CalledProcessError as e:
335
+ self.log_error(f"Update failed (Exit Code {e.returncode}):\nSTDOUT:\n{e.stdout}\nSTDERR:\n{e.stderr}")
336
+ return f"failed_exit_{e.returncode}"
337
+ except subprocess.TimeoutExpired:
338
+ self.log_error(f"Update command timed out.")
339
+ return "failed_timeout"
340
+ except Exception as e:
341
+ self.log_error(f"Unexpected error during update: {e}", exc_info=True)
342
+ return "failed_exception"
343
+
344
+ def _get_ootb_processes(self, target_user="all_active"):
345
+ ootb_procs = []
346
+ target_pid_list = []
347
+ try:
348
+ target_users = set()
349
+ if target_user == "all_active":
350
+ for user_session in psutil.users():
351
+ username = user_session.name.split('\\')[-1]
352
+ target_users.add(username.lower())
353
+ else:
354
+ target_users.add(target_user.lower())
355
+ self.log_info(f"Searching for OOTB processes for users: {target_users}")
356
+
357
+ # Use the potentially corrected python.exe path for matching
358
+ python_exe_path_for_check = self.target_executable_path.strip('"')
359
+ self.log_info(f"_get_ootb_processes: Checking against python path: {python_exe_path_for_check}")
360
+
361
+ for proc in psutil.process_iter(['pid', 'name', 'username', 'cmdline', 'exe']):
362
+ try:
363
+ pinfo = proc.info
364
+ proc_username = pinfo['username']
365
+ if proc_username:
366
+ proc_username = proc_username.split('\\')[-1].lower()
367
+
368
+ if proc_username in target_users:
369
+ cmdline = ' '.join(pinfo['cmdline']) if pinfo['cmdline'] else ''
370
+ # Check if the process executable matches our corrected python path AND module is in cmdline
371
+ if pinfo['exe'] and pinfo['exe'] == python_exe_path_for_check and _OOTB_MODULE in cmdline:
372
+ self.log_info(f"Found matching OOTB process: PID={pinfo['pid']}, User={pinfo['username']}, Cmd={cmdline}")
373
+ ootb_procs.append(proc)
374
+ target_pid_list.append(pinfo['pid'])
375
+ except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
376
+ continue
377
+ self.log_info(f"Found {len(ootb_procs)} OOTB process(es) matching criteria: {target_pid_list}")
378
+ except Exception as e:
379
+ self.log_error(f"Error enumerating processes: {e}", exc_info=True)
380
+ return ootb_procs
381
+
382
+ def handle_stop(self, target_user="all_active"):
383
+ self.log_info(f"Executing stop OOTB for target '{target_user}'...")
384
+ stop_results = {} # Track results per user {username: (task_status, immediate_status)}
385
+ failed_users = set()
386
+
387
+ try:
388
+ # --- Get target users and active sessions ---
389
+ active_sessions = {} # user_lower: session_id
390
+ # No need for all_system_users for stop, we only care about active or the specific target
391
+ try:
392
+ sessions = win32ts.WTSEnumerateSessions(win32ts.WTS_CURRENT_SERVER_HANDLE)
393
+ for session in sessions:
394
+ if session['State'] == win32ts.WTSActive:
395
+ try:
396
+ user = win32ts.WTSQuerySessionInformation(win32ts.WTS_CURRENT_SERVER_HANDLE, session['SessionId'], win32ts.WTSUserName)
397
+ if user:
398
+ active_sessions[user.lower()] = session['SessionId']
399
+ except Exception as query_err:
400
+ self.log_error(f"Could not query session {session['SessionId']} during stop: {query_err}")
401
+ except Exception as user_enum_err:
402
+ self.log_error(f"Error enumerating users/sessions during stop: {user_enum_err}", exc_info=True)
403
+ return "failed_user_enumeration"
404
+
405
+ self.log_info(f"Stop target: '{target_user}'. Active sessions: {active_sessions}")
406
+
407
+ target_users_normalized = set()
408
+ if target_user == "all_active":
409
+ # Target only currently active users for stop all
410
+ target_users_normalized = set(active_sessions.keys())
411
+ self.log_info(f"Stop targeting all active users: {target_users_normalized}")
412
+ else:
413
+ # Target the specific user, regardless of active status (for task removal)
414
+ normalized_target = target_user.lower()
415
+ target_users_normalized.add(normalized_target)
416
+ self.log_info(f"Stop targeting specific user: {normalized_target}")
417
+
418
+ if not target_users_normalized:
419
+ self.log_info("No target users identified for stop.")
420
+ return "failed_no_target_users" # Or success if none were targeted?
421
+
422
+ # --- Process each target user ---
423
+ for user in target_users_normalized:
424
+ task_removed_status = "task_unknown"
425
+ immediate_stop_status = "stop_not_attempted"
426
+ stopped_count = 0
427
+
428
+ self.log_info(f"Processing stop for user '{user}'...")
429
+
430
+ # 1. Always try to remove the scheduled task
431
+ try:
432
+ # remove_logon_task always returns True for now, just logs attempt
433
+ self.remove_logon_task(user)
434
+ task_removed_status = "task_removed_attempted"
435
+ except Exception as task_err:
436
+ self.log_error(f"Exception removing scheduled task for {user}: {task_err}", exc_info=True)
437
+ task_removed_status = "task_exception"
438
+ failed_users.add(user)
439
+ # Continue to try and stop process if active
440
+
441
+ # 2. If user is active, try to terminate process
442
+ is_active = user in active_sessions
443
+
444
+ if is_active:
445
+ immediate_stop_status = "stop_attempted"
446
+ self.log_info(f"User '{user}' is active. Attempting to terminate OOTB process(es)...")
447
+ # Pass the specific username to _get_ootb_processes
448
+ procs_to_stop = self._get_ootb_processes(user)
449
+
450
+ if not procs_to_stop:
451
+ self.log_info(f"No running OOTB processes found for active user '{user}'.")
452
+ immediate_stop_status = "stop_skipped_not_running"
453
+ else:
454
+ self.log_info(f"Found {len(procs_to_stop)} process(es) for user '{user}' to stop.")
455
+ for proc in procs_to_stop:
456
+ try:
457
+ pid = proc.pid # Get pid before potential termination
458
+ username = proc.info.get('username', 'unknown_user')
459
+ self.log_info(f"Terminating process PID={pid}, User={username}")
460
+ proc.terminate()
461
+ try:
462
+ proc.wait(timeout=3)
463
+ self.log_info(f"Process PID={pid} terminated successfully.")
464
+ stopped_count += 1
465
+ except psutil.TimeoutExpired:
466
+ self.log_error(f"Process PID={pid} did not terminate gracefully, killing.")
467
+ proc.kill()
468
+ stopped_count += 1
469
+ except psutil.NoSuchProcess:
470
+ self.log_info(f"Process PID={pid} already terminated.")
471
+ # Don't increment stopped_count here as we didn't stop it now
472
+ except psutil.AccessDenied:
473
+ self.log_error(f"Access denied trying to terminate process PID={pid}.")
474
+ failed_users.add(user) # Mark user as failed if stop fails
475
+ except Exception as e:
476
+ self.log_error(f"Error stopping process PID={pid}: {e}", exc_info=True)
477
+ failed_users.add(user) # Mark user as failed
478
+
479
+ # Determine status based on how many were found vs stopped
480
+ if user in failed_users:
481
+ immediate_stop_status = f"stop_errors_terminated_{stopped_count}_of_{len(procs_to_stop)}"
482
+ elif stopped_count == len(procs_to_stop):
483
+ immediate_stop_status = f"stop_success_terminated_{stopped_count}"
484
+ else: # Should ideally not happen if NoSuchProcess doesn't count
485
+ immediate_stop_status = f"stop_partial_terminated_{stopped_count}_of_{len(procs_to_stop)}"
486
+
487
+ else: # User not active
488
+ self.log_info(f"User '{user}' is not active. Skipping immediate process stop (task removal attempted).")
489
+ immediate_stop_status = "stop_skipped_inactive"
490
+
491
+ # Record final results for this user
492
+ stop_results[user] = (task_removed_status, immediate_stop_status)
493
+
494
+
495
+ # --- Consolidate status ---
496
+ total_processed = len(target_users_normalized)
497
+ final_status = "partial_success" if failed_users else "success"
498
+ if not stop_results: final_status = "no_targets_processed"
499
+ if len(failed_users) == total_processed and total_processed > 0 : final_status = "failed"
500
+
501
+ self.log_info(f"Finished stopping OOTB. Overall Status: {final_status}. Results: {stop_results}")
502
+ try:
503
+ details = json.dumps(stop_results)
504
+ except Exception:
505
+ details = str(stop_results) # Fallback
506
+ return f"{final_status}::{details}" # Use :: as separator
507
+
508
+ except Exception as e:
509
+ self.log_error(f"Error during combined stop OOTB process: {e}", exc_info=True)
510
+ return "failed_exception"
511
+
512
+
513
+ def handle_start(self, target_user="all_active"):
514
+ self.log_info(f"Executing start OOTB for target '{target_user}'...")
515
+ start_results = {} # Track results per user {username: (task_status, immediate_status)}
516
+ failed_users = set()
517
+
518
+ try:
519
+ # --- Get target users and active sessions ---
520
+ active_sessions = {} # user_lower: session_id
521
+ all_system_users = set() # user_lower
522
+ try:
523
+ # Use psutil for system user list, WTS for active sessions/IDs
524
+ for user_session in psutil.users():
525
+ username_lower = user_session.name.split('\\')[-1].lower()
526
+ all_system_users.add(username_lower)
527
+
528
+ sessions = win32ts.WTSEnumerateSessions(win32ts.WTS_CURRENT_SERVER_HANDLE)
529
+ for session in sessions:
530
+ if session['State'] == win32ts.WTSActive:
531
+ try:
532
+ user = win32ts.WTSQuerySessionInformation(win32ts.WTS_CURRENT_SERVER_HANDLE, session['SessionId'], win32ts.WTSUserName)
533
+ if user:
534
+ active_sessions[user.lower()] = session['SessionId']
535
+ except Exception as query_err:
536
+ self.log_error(f"Could not query session {session['SessionId']}: {query_err}")
537
+ except Exception as user_enum_err:
538
+ self.log_error(f"Error enumerating users/sessions: {user_enum_err}", exc_info=True)
539
+ return "failed_user_enumeration"
540
+
541
+ self.log_info(f"Found active user sessions: {active_sessions}")
542
+
543
+ target_users_normalized = set()
544
+ if target_user == "all_active":
545
+ # If targeting all_active, only target those CURRENTLY active
546
+ target_users_normalized = set(active_sessions.keys())
547
+ self.log_info(f"Targeting all active users: {target_users_normalized}")
548
+ else:
549
+ normalized_target = target_user.lower()
550
+ # Check if the target user actually exists on the system, even if inactive
551
+ # This check might be complex/unreliable. Rely on task scheduler potentially failing?
552
+ # Let's assume admin provides a valid username for specific targeting.
553
+ # if normalized_target in all_system_users: # Removing this check, assume valid user input
554
+ target_users_normalized.add(normalized_target)
555
+ self.log_info(f"Targeting specific user: {normalized_target}")
556
+ # else:
557
+ # log_error(f"Target user '{target_user}' does not appear to exist on this system based on psutil.")
558
+ # return "failed_user_does_not_exist"
559
+
560
+ if not target_users_normalized:
561
+ self.log_info("No target users identified (or none active for 'all_active').")
562
+ # If target was specific user but they weren't found active, still try task?
563
+ # Let's proceed to task creation anyway for specific user case.
564
+ if target_user != "all_active": target_users_normalized.add(target_user.lower())
565
+ if not target_users_normalized:
566
+ return "failed_no_target_users"
567
+
568
+ # --- Check existing processes ---
569
+ # This check is only relevant for immediate start attempt
570
+ running_procs_by_user = {} # user_lower: count
571
+ try:
572
+ current_running = self._get_ootb_processes("all_active") # Check all
573
+ for proc in current_running:
574
+ try:
575
+ proc_username = proc.info.get('username')
576
+ if proc_username:
577
+ user_lower = proc_username.split('\\')[-1].lower()
578
+ running_procs_by_user[user_lower] = running_procs_by_user.get(user_lower, 0) + 1
579
+ except Exception: pass
580
+ except Exception as e:
581
+ self.log_error(f"Error checking existing processes: {e}")
582
+ self.log_info(f"Users currently running OOTB: {running_procs_by_user}")
583
+
584
+ # --- Process each target user ---
585
+ for user in target_users_normalized:
586
+ task_created_status = "task_unknown"
587
+ immediate_start_status = "start_not_attempted"
588
+ token = None # Ensure token is reset/defined
589
+
590
+ self.log_info(f"Processing start for user '{user}'...")
591
+
592
+ # 1. Always try to create/update the scheduled task
593
+ try:
594
+ task_created = self.create_or_update_logon_task(user)
595
+ task_created_status = "task_success" if task_created else "task_failed"
596
+ except Exception as task_err:
597
+ self.log_error(f"Exception creating/updating scheduled task for {user}: {task_err}", exc_info=True)
598
+ task_created_status = "task_exception"
599
+ failed_users.add(user)
600
+ # Continue to potentially try immediate start IF user is active?
601
+ # Or maybe skip if task creation failed badly?
602
+ # Let's skip immediate start if task creation had exception.
603
+ start_results[user] = (task_created_status, immediate_start_status)
604
+ continue
605
+
606
+ # 2. If user is active AND not already running, try immediate start
607
+ is_active = user in active_sessions
608
+ is_running = running_procs_by_user.get(user, 0) > 0
609
+
610
+ if is_active:
611
+ if not is_running:
612
+ immediate_start_status = "start_attempted"
613
+ self.log_info(f"User '{user}' is active and not running OOTB. Attempting immediate start...")
614
+ try:
615
+ session_id = active_sessions[user]
616
+ token = win32ts.WTSQueryUserToken(session_id)
617
+ env = win32profile.CreateEnvironmentBlock(token, False) # <-- Restore creating env block
618
+ startup = win32process.STARTUPINFO()
619
+ creation_flags = 0x00000010 # CREATE_NEW_CONSOLE
620
+
621
+ # --- Launch via cmd /K ---
622
+ lpApplicationName = None
623
+ lpCommandLine = f'cmd.exe /K "{self.target_executable_path}"'
624
+ cwd = os.path.dirname(self.target_executable_path.strip('"')) if os.path.dirname(self.target_executable_path.strip('"')) != '' else None
625
+ # --- End Launch via cmd /K ---
626
+
627
+ self.log_info(f"Calling CreateProcessAsUser to launch via cmd /K (Restored Environment Block):") # Updated log note
628
+ self.log_info(f" lpApplicationName: {lpApplicationName}")
629
+ self.log_info(f" lpCommandLine: {lpCommandLine}")
630
+ self.log_info(f" lpCurrentDirectory: {cwd if cwd else 'Default'}")
631
+ self.log_info(f" dwCreationFlags: {creation_flags} (CREATE_NEW_CONSOLE)")
632
+
633
+ # CreateProcessAsUser call with restored env
634
+ hProcess, hThread, dwPid, dwTid = win32process.CreateProcessAsUser(
635
+ token, # User token
636
+ lpApplicationName, # Application name (None)
637
+ lpCommandLine, # Command line (cmd.exe /K "...")
638
+ None, # Process attributes
639
+ None, # Thread attributes
640
+ False, # Inherit handles
641
+ creation_flags, # Creation flags (CREATE_NEW_CONSOLE)
642
+ env, # Environment block (Restored)
643
+ cwd, # Current directory for cmd.exe
644
+ startup # Startup info
645
+ )
646
+
647
+ self.log_info(f"CreateProcessAsUser call succeeded for user '{user}' (PID: {dwPid}). Checking existence...")
648
+ win32api.CloseHandle(hProcess)
649
+ win32api.CloseHandle(hThread)
650
+
651
+ time.sleep(1)
652
+ if psutil.pid_exists(dwPid):
653
+ self.log_info(f"Immediate start succeeded for user '{user}' (PID {dwPid}).")
654
+ immediate_start_status = "start_success"
655
+ else:
656
+ self.log_error(f"Immediate start failed for user '{user}': Process {dwPid} exited immediately.")
657
+ immediate_start_status = "start_failed_exited"
658
+ failed_users.add(user)
659
+
660
+ except Exception as proc_err:
661
+ self.log_error(f"Exception during immediate start for user '{user}': {proc_err}", exc_info=True)
662
+ immediate_start_status = "start_failed_exception"
663
+ failed_users.add(user)
664
+ finally:
665
+ if token:
666
+ try: win32api.CloseHandle(token)
667
+ except: pass
668
+ else: # User is active but already running
669
+ self.log_info(f"User '{user}' is active but OOTB is already running. Skipping immediate start.")
670
+ immediate_start_status = "start_skipped_already_running"
671
+ else: # User is not active
672
+ self.log_info(f"User '{user}' is not active. Skipping immediate start (task created/updated).")
673
+ immediate_start_status = "start_skipped_inactive"
674
+
675
+ # Record final results for this user
676
+ start_results[user] = (task_created_status, immediate_start_status)
677
+
678
+
679
+ # --- Consolidate status ---
680
+ total_processed = len(target_users_normalized)
681
+ final_status = "partial_success" if failed_users else "success"
682
+ if not start_results: final_status = "no_targets_processed"
683
+ # If all processed users failed in some way (either task or start)
684
+ if len(failed_users) == total_processed and total_processed > 0: final_status = "failed"
685
+ # Special case: target was specific user who wasn't found active
686
+ elif total_processed == 1 and target_user != "all_active" and target_user.lower() not in active_sessions:
687
+ user_key = target_user.lower()
688
+ if user_key in start_results and start_results[user_key][0] == "task_success":
689
+ final_status = "success_task_only_user_inactive"
690
+ else:
691
+ final_status = "failed_task_user_inactive"
692
+
693
+ self.log_info(f"Finished starting OOTB. Overall Status: {final_status}. Results: {start_results}")
694
+ # Return detailed results as a JSON string for easier parsing/logging server-side
695
+ try:
696
+ details = json.dumps(start_results)
697
+ except Exception:
698
+ details = str(start_results) # Fallback
699
+ return f"{final_status}::{details}"
700
+
701
+ except Exception as e:
702
+ self.log_error(f"Error during combined start OOTB process: {e}", exc_info=True)
703
+ return "failed_exception"
704
+
705
+ def create_or_update_logon_task(self, username):
706
+ """Creates/updates task to run OOTB app via cmd /K on session connect/reconnect."""
707
+ if not self.target_executable_path:
708
+ self.log_error(f"Cannot create task for {username}: Target executable path is not set.")
709
+ return False
710
+
711
+ task_name = f"OOTB_UserConnect_{username}"
712
+ # Action: Revert to running the actual executable via cmd /K
713
+ action_executable = 'cmd.exe'
714
+ action_arguments = f'/K "{self.target_executable_path}"'
715
+ safe_action_executable = action_executable.replace("'", "''")
716
+ safe_action_arguments = action_arguments.replace("'", "''")
717
+
718
+ # Explicitly set the working directory to the executable's location
719
+ try:
720
+ executable_dir = os.path.dirname(self.target_executable_path.strip('"'))
721
+ if not executable_dir: executable_dir = "."
722
+ safe_working_directory = executable_dir.replace("'", "''")
723
+ working_directory_setting = f"$action.WorkingDirectory = '{safe_working_directory}'"
724
+ except Exception as e:
725
+ self.log_error(f"Error determining working directory for task: {e}. WD will not be set.")
726
+ working_directory_setting = "# Could not set WorkingDirectory"
727
+
728
+ # PowerShell command construction
729
+ ps_command = f"""
730
+ $taskName = "{task_name}"
731
+ $principal = New-ScheduledTaskPrincipal -UserId "{username}" -LogonType Interactive
732
+
733
+ # Action: Run the OOTB executable via cmd /K
734
+ $action = New-ScheduledTaskAction -Execute '{safe_action_executable}' -Argument '{safe_action_arguments}'
735
+ {working_directory_setting} # Set the working directory
736
+
737
+ # Triggers: On session connect (21) AND reconnect (25)
738
+ $logName = 'Microsoft-Windows-TerminalServices-LocalSessionManager/Operational'
739
+ $source = 'Microsoft-Windows-TerminalServices-LocalSessionManager'
740
+ # Define multiple triggers
741
+ $trigger1 = New-ScheduledTaskTrigger -Event -LogName $logName -Source $source -EventId 21
742
+ $trigger2 = New-ScheduledTaskTrigger -Event -LogName $logName -Source $source -EventId 25
743
+ # Optional Delay - Apply to both triggers if desired?
744
+ # $trigger1.Delay = 'PT15S'
745
+ # $trigger2.Delay = 'PT15S'
746
+
747
+ $settings = New-ScheduledTaskSettingsSet -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries -ExecutionTimeLimit (New-TimeSpan -Days 9999) -RestartCount 3 -RestartInterval (New-TimeSpan -Minutes 1)
748
+ $description = "Runs OOTB Application (via cmd) for user {username} upon session connect/reconnect." # Updated description
749
+
750
+ # Unregister existing task first (force) - Use the NEW task name
751
+ Unregister-ScheduledTask -TaskName $taskName -Confirm:$false -ErrorAction SilentlyContinue
752
+
753
+ # Register the new task with MULTIPLE triggers
754
+ Register-ScheduledTask -TaskName $taskName -Principal $principal -Action $action -Trigger $trigger1, $trigger2 -Settings $settings -Description $description -Force
755
+
756
+ """
757
+ self.log_info(f"Attempting to create/update task '{task_name}' for user '{username}' to run OOTB on session connect/reconnect.")
758
+ try:
759
+ # Need to actually run the powershell command here!
760
+ success = self.run_powershell_command(ps_command)
761
+ if success:
762
+ self.log_info(f"Successfully ran PowerShell command to create/update task '{task_name}'.")
763
+ return True
764
+ else:
765
+ self.log_error(f"PowerShell command failed to create/update task '{task_name}'. See previous logs.")
766
+ return False
767
+ except Exception as e:
768
+ self.log_error(f"Failed to create/update scheduled task '{task_name}' for user '{username}': {e}", exc_info=True)
769
+ return False
770
+
771
+ def run_powershell_command(self, command, log_output=True):
772
+ """Executes a PowerShell command and handles output/errors. Returns True on success."""
773
+ self.log_info(f"Executing PowerShell: {command}")
774
+ try:
775
+ result = subprocess.run(
776
+ ["powershell.exe", "-NoProfile", "-NonInteractive", "-Command", command],
777
+ capture_output=True, text=True, check=True, encoding='utf-8', errors='ignore'
778
+ )
779
+ if log_output and result.stdout:
780
+ self.log_info(f"PowerShell STDOUT:\n{result.stdout.strip()}")
781
+ if log_output and result.stderr:
782
+ # Log stderr as info, as some commands write status here (like unregister task not found)
783
+ self.log_info(f"PowerShell STDERR:\n{result.stderr.strip()}")
784
+ return True
785
+ except FileNotFoundError:
786
+ self.log_error("'powershell.exe' not found. Cannot manage scheduled tasks.")
787
+ return False
788
+ except subprocess.CalledProcessError as e:
789
+ # Log error but still return False, handled by caller
790
+ self.log_error(f"PowerShell command failed (Exit Code {e.returncode}):")
791
+ self.log_error(f" Command: {e.cmd}")
792
+ if e.stdout: self.log_error(f" STDOUT: {e.stdout.strip()}")
793
+ if e.stderr: self.log_error(f" STDERR: {e.stderr.strip()}")
794
+ return False
795
+ except Exception as e:
796
+ self.log_error(f"Unexpected error running PowerShell: {e}", exc_info=True)
797
+ return False
798
+
799
+ def remove_logon_task(self, username):
800
+ """Removes the logon scheduled task for a user."""
801
+ task_name = f"{self._task_name_prefix}{username}"
802
+ safe_task_name = task_name.replace("'", "''")
803
+ command = f"Unregister-ScheduledTask -TaskName '{safe_task_name}' -Confirm:$false -ErrorAction SilentlyContinue"
804
+ self.run_powershell_command(command, log_output=False)
805
+ self.log_info(f"Attempted removal of scheduled task '{task_name}' for user '{username}'.")
806
+ return True
807
+
808
+ # --- Main Execution Block ---
809
+ if __name__ == '__main__':
810
+ if len(sys.argv) > 1 and sys.argv[1] == 'debug':
811
+ self.log_info("Starting service in debug mode...")
812
+ print(f"Running Flask server via Waitress on {_LISTEN_HOST}:{_LISTEN_PORT} for debugging...")
813
+ print("Service logic (command processing) will NOT run in this mode.")
814
+ print("Use this primarily to test the '/command' endpoint receiving POSTs.")
815
+ print("Press Ctrl+C to stop.")
816
+ try:
817
+ serve(flask_app, host=_LISTEN_HOST, port=_LISTEN_PORT, threads=1)
818
+ except KeyboardInterrupt:
819
+ print("\nDebug server stopped.")
820
+
821
+ elif len(sys.argv) == 1:
822
+ try:
823
+ servicemanager.Initialize()
824
+ servicemanager.PrepareToHostSingle(GuardService)
825
+ servicemanager.StartServiceCtrlDispatcher()
826
+ except win32service.error as details:
827
+ import winerror
828
+ if details.winerror == winerror.ERROR_FAILED_SERVICE_CONTROLLER_CONNECT:
829
+ print(f"Error: Not started by SCM.")
830
+ print(f"Use 'python {os.path.basename(__file__)} install|start|stop|remove|debug'")
831
+ else:
832
+ print(f"Error preparing service: {details}")
833
+ except Exception as e:
834
+ print(f"Unexpected error initializing service: {e}")
835
+ else:
836
836
  win32serviceutil.HandleCommandLine(GuardService)