windows-ui-automation

from comtypes.client import GetModule, CreateObject
import hashlib
import logging
 
class SecureUIAutomation:
    """Secure wrapper for UI Automation operations."""
 
    BLOCKED_PROCESSES = {
        'keepass.exe', '1password.exe', 'lastpass.exe',    # Password managers
        'mmc.exe', 'secpol.msc', 'gpedit.msc',             # Admin tools
        'regedit.exe', 'cmd.exe', 'powershell.exe',        # System tools
        'taskmgr.exe', 'procexp.exe',                       # Process tools
    }
 
    def __init__(self, permission_tier: str = 'read-only'):
        self.permission_tier = permission_tier
        self.uia = CreateObject('UIAutomationClient.CUIAutomation')
        self.logger = logging.getLogger('uia.security')
        self.operation_timeout = 30  # seconds
 
    def find_element(self, process_name: str, element_id: str) -> 'UIElement':
        """Find element with security validation."""
        # Security check: blocked processes
        if process_name.lower() in self.BLOCKED_PROCESSES:
            self.logger.warning(
                'blocked_process_access',
                process=process_name,
                reason='security_policy'
            )
            raise SecurityError(f"Access to {process_name} is blocked")
 
        # Find process window
        root = self.uia.GetRootElement()
        condition = self.uia.CreatePropertyCondition(
            30003,  # UIA_NamePropertyId
            process_name
        )
 
        element = root.FindFirst(4, condition)  # TreeScope_Children
 
        if element:
            self._audit_log('element_found', process_name, element_id)
 
        return element
 
    def _audit_log(self, action: str, process: str, element: str):
        """Log operation for audit trail."""
        self.logger.info(
            f'uia.{action}',
            extra={
                'process': process,
                'element': element,
                'permission_tier': self.permission_tier,
                'correlation_id': self._get_correlation_id()
            }
        )

import ctypes
from ctypes import wintypes
import time
 
class SafeInputSimulator:
    """Input simulation with security controls."""
 
    # Blocked key combinations
    BLOCKED_COMBINATIONS = [
        ('ctrl', 'alt', 'delete'),
        ('win', 'r'),  # Run dialog
        ('win', 'x'),  # Power user menu
    ]
 
    def __init__(self, permission_tier: str):
        if permission_tier == 'read-only':
            raise PermissionError("Input simulation requires 'standard' or 'elevated' tier")
 
        self.permission_tier = permission_tier
        self.rate_limit = 100  # max inputs per second
        self._input_count = 0
        self._last_reset = time.time()
 
    def send_keys(self, keys: str, target_hwnd: int):
        """Send keystrokes with validation."""
        # Rate limiting
        self._check_rate_limit()
 
        # Validate target window
        if not self._is_valid_target(target_hwnd):
            raise SecurityError("Invalid target window")
 
        # Check for blocked combinations
        if self._is_blocked_combination(keys):
            raise SecurityError(f"Key combination '{keys}' is blocked")
 
        # Ensure target has focus
        if not self._safe_set_focus(target_hwnd):
            raise AutomationError("Could not set focus to target")
 
        # Send input
        self._send_input_safe(keys)
 
    def _check_rate_limit(self):
        """Prevent input flooding."""
        now = time.time()
        if now - self._last_reset > 1.0:
            self._input_count = 0
            self._last_reset = now
 
        self._input_count += 1
        if self._input_count > self.rate_limit:
            raise RateLimitError("Input rate limit exceeded")

import psutil
import hashlib
 
class ProcessValidator:
    """Validate processes before automation."""
 
    def __init__(self):
        self.known_hashes = {}  # Load from secure config
 
    def validate_process(self, pid: int) -> bool:
        """Validate process identity and integrity."""
        try:
            proc = psutil.Process(pid)
 
            # Check process name against blocklist
            if proc.name().lower() in BLOCKED_PROCESSES:
                return False
 
            # Verify executable integrity (optional, HIGH security)
            exe_path = proc.exe()
            if not self._verify_integrity(exe_path):
                return False
 
            # Check process owner
            if not self._check_owner(proc):
                return False
 
            return True
 
        except psutil.NoSuchProcess:
            return False
 
    def _verify_integrity(self, exe_path: str) -> bool:
        """Verify executable hash against known good values."""
        if exe_path not in self.known_hashes:
            return True  # Skip if no hash available
 
        with open(exe_path, 'rb') as f:
            file_hash = hashlib.sha256(f.read()).hexdigest()
 
        return file_hash == self.known_hashes[exe_path]

import signal
from contextlib import contextmanager
 
class TimeoutManager:
    """Enforce operation timeouts."""
 
    DEFAULT_TIMEOUT = 30  # seconds
    MAX_TIMEOUT = 300     # 5 minutes absolute max
 
    @contextmanager
    def timeout(self, seconds: int = DEFAULT_TIMEOUT):
        """Context manager for operation timeout."""
        if seconds > self.MAX_TIMEOUT:
            seconds = self.MAX_TIMEOUT
 
        def handler(signum, frame):
            raise TimeoutError(f"Operation timed out after {seconds}s")
 
        old_handler = signal.signal(signal.SIGALRM, handler)
        signal.alarm(seconds)
 
        try:
            yield
        finally:
            signal.alarm(0)
            signal.signal(signal.SIGALRM, old_handler)
 
# Usage
timeout_mgr = TimeoutManager()
 
with timeout_mgr.timeout(10):
    element = automation.find_element('notepad.exe', 'Edit1')

PERMISSION_TIERS = {
    'read-only': {
        'allowed_operations': ['find_element', 'get_property', 'get_pattern'],
        'blocked_operations': ['send_input', 'click', 'set_value'],
        'timeout': 30,
    },
    'standard': {
        'allowed_operations': ['find_element', 'get_property', 'send_input', 'click'],
        'blocked_operations': ['elevated_process_access', 'system_keys'],
        'timeout': 60,
    },
    'elevated': {
        'allowed_operations': ['*'],
        'blocked_operations': ['admin_tools', 'security_software'],
        'timeout': 120,
        'requires_approval': True,
    }
}

# tests/test_ui_automation.py
import pytest
from unittest.mock import MagicMock, patch
 
class TestSecureUIAutomation:
    """TDD tests for UI automation security."""
 
    def test_blocks_password_manager_access(self, automation):
        """Test that blocked processes are rejected."""
        with pytest.raises(SecurityError, match="blocked"):
            automation.find_element('keepass.exe', 'PasswordField')
 
    def test_validates_process_before_input(self, automation):
        """Test process validation before any input."""
        with patch.object(automation, '_validate_process') as mock_validate:
            mock_validate.return_value = False
            with pytest.raises(SecurityError):
                automation.send_keys('test', hwnd=12345)
            mock_validate.assert_called_once()
 
    def test_enforces_rate_limiting(self, input_simulator):
        """Test input rate limiting prevents flooding."""
        for _ in range(100):
            input_simulator.send_keys('a', hwnd=12345)
        with pytest.raises(RateLimitError):
            input_simulator.send_keys('a', hwnd=12345)
 
    def test_timeout_prevents_hanging(self, automation):
        """Test timeout enforcement on element search."""
        with pytest.raises(TimeoutError):
            with automation.timeout(0.001):
                automation.find_element('app.exe', 'NonExistent')
 
@pytest.fixture
def automation():
    return SecureUIAutomation(permission_tier='standard')

class SecureUIAutomation:
    BLOCKED_PROCESSES = {'keepass.exe', '1password.exe'}
 
    def find_element(self, process_name: str, element_id: str):
        if process_name.lower() in self.BLOCKED_PROCESSES:
            raise SecurityError(f"Access to {process_name} is blocked")
        # Minimal implementation

# Run all tests with coverage
pytest tests/test_ui_automation.py -v --cov=src/automation --cov-report=term-missing
 
# Run security-specific tests
pytest tests/ -k "security or blocked" -v
 
# Type checking
mypy src/automation --strict

# BAD: Re-find element every operation
for i in range(100):
    element = uia.find_element('app.exe', 'TextField')
    element.send_keys(str(i))
 
# GOOD: Cache element reference
element = uia.find_element('app.exe', 'TextField')
for i in range(100):
    if element.is_valid():
        element.send_keys(str(i))
    else:
        element = uia.find_element('app.exe', 'TextField')

# BAD: Search from root every time
root = uia.GetRootElement()
element = root.FindFirst(TreeScope.Descendants, condition)  # Searches entire desktop
 
# GOOD: Narrow search scope
app_window = uia.find_window('notepad.exe')
element = app_window.FindFirst(TreeScope.Children, condition)  # Only direct children

# BAD: Blocking wait for element
while not element.is_enabled():
    time.sleep(0.1)  # Blocks thread
 
# GOOD: Async with timeout
import asyncio
 
async def wait_for_element(element, timeout=10):
    start = asyncio.get_event_loop().time()
    while not element.is_enabled():
        if asyncio.get_event_loop().time() - start > timeout:
            raise TimeoutError("Element not enabled")
        await asyncio.sleep(0.05)  # Non-blocking

# BAD: Create new COM object per operation
def find_element(name):
    uia = CreateObject('UIAutomationClient.CUIAutomation')  # Expensive
    return uia.GetRootElement().FindFirst(...)
 
# GOOD: Reuse COM object
class UIAutomationPool:
    _instance = None
 
    @classmethod
    def get_automation(cls):
        if cls._instance is None:
            cls._instance = CreateObject('UIAutomationClient.CUIAutomation')
        return cls._instance

# BAD: Multiple sequential conditions
name_cond = uia.CreatePropertyCondition(UIA_NamePropertyId, 'Submit')
type_cond = uia.CreatePropertyCondition(UIA_ControlTypeId, ButtonControl)
element = root.FindFirst(TreeScope.Descendants, name_cond)
if element.ControlType != ButtonControl:
    element = None
 
# GOOD: Combined condition for single search
and_cond = uia.CreateAndCondition(
    uia.CreatePropertyCondition(UIA_NamePropertyId, 'Submit'),
    uia.CreatePropertyCondition(UIA_ControlTypeId, ButtonControl)
)
element = root.FindFirst(TreeScope.Descendants, and_cond)

# BAD: No validation
element = uia.find_element_by_name('Password')
element.send_keys(password)
 
# GOOD: Full validation
if validator.validate_process(target_pid):
    if automation.permission_tier != 'read-only':
        element = automation.find_element(process_name, 'Password')
        element.send_keys(password)

# BAD: No timeout
element = uia.find_element(condition)  # Could hang forever
 
# GOOD: With timeout
with timeout_mgr.timeout(10):
    element = uia.find_element(condition)

# BAD: Allow any keys
def send_keys(keys):
    SendInput(keys)
 
# GOOD: Block dangerous combinations
def send_keys(keys):
    if is_blocked_combination(keys):
        raise SecurityError("Blocked key combination")
    SendInput(keys)