Initial commit
This commit is contained in:
Zhongwei Li
628
scripts/utils/helpers.py
Normal file
628
scripts/utils/helpers.py
Normal file
@@ -0,0 +1,628 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Helper utilities for Tailscale SSH Sync Agent.
|
||||
Provides common formatting, parsing, and utility functions.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, Any
|
||||
import yaml
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def format_bytes(bytes_value: int) -> str:
|
||||
"""
|
||||
Format bytes as human-readable string.
|
||||
|
||||
Args:
|
||||
bytes_value: Number of bytes
|
||||
|
||||
Returns:
|
||||
Formatted string (e.g., "12.3 MB", "1.5 GB")
|
||||
|
||||
Example:
|
||||
>>> format_bytes(12582912)
|
||||
"12.0 MB"
|
||||
>>> format_bytes(1610612736)
|
||||
"1.5 GB"
|
||||
"""
|
||||
for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
|
||||
if bytes_value < 1024.0:
|
||||
return f"{bytes_value:.1f} {unit}"
|
||||
bytes_value /= 1024.0
|
||||
return f"{bytes_value:.1f} PB"
|
||||
|
||||
|
||||
def format_duration(seconds: float) -> str:
|
||||
"""
|
||||
Format duration as human-readable string.
|
||||
|
||||
Args:
|
||||
seconds: Duration in seconds
|
||||
|
||||
Returns:
|
||||
Formatted string (e.g., "2m 15s", "1h 30m")
|
||||
|
||||
Example:
|
||||
>>> format_duration(135)
|
||||
"2m 15s"
|
||||
>>> format_duration(5430)
|
||||
"1h 30m 30s"
|
||||
"""
|
||||
if seconds < 60:
|
||||
return f"{int(seconds)}s"
|
||||
|
||||
minutes = int(seconds // 60)
|
||||
secs = int(seconds % 60)
|
||||
|
||||
if minutes < 60:
|
||||
return f"{minutes}m {secs}s" if secs > 0 else f"{minutes}m"
|
||||
|
||||
hours = minutes // 60
|
||||
minutes = minutes % 60
|
||||
|
||||
parts = [f"{hours}h"]
|
||||
if minutes > 0:
|
||||
parts.append(f"{minutes}m")
|
||||
if secs > 0 and hours == 0: # Only show seconds if < 1 hour
|
||||
parts.append(f"{secs}s")
|
||||
|
||||
return " ".join(parts)
|
||||
|
||||
|
||||
def format_percentage(value: float, decimals: int = 1) -> str:
|
||||
"""
|
||||
Format percentage with specified decimals.
|
||||
|
||||
Args:
|
||||
value: Percentage value (0-100)
|
||||
decimals: Number of decimal places
|
||||
|
||||
Returns:
|
||||
Formatted string (e.g., "45.5%")
|
||||
|
||||
Example:
|
||||
>>> format_percentage(45.567)
|
||||
"45.6%"
|
||||
"""
|
||||
return f"{value:.{decimals}f}%"
|
||||
|
||||
|
||||
def parse_ssh_config(config_path: Optional[Path] = None) -> Dict[str, Dict[str, str]]:
|
||||
"""
|
||||
Parse SSH config file for host definitions.
|
||||
|
||||
Args:
|
||||
config_path: Path to SSH config (default: ~/.ssh/config)
|
||||
|
||||
Returns:
|
||||
Dict mapping host aliases to their configuration:
|
||||
{
|
||||
'host-alias': {
|
||||
'hostname': '100.64.1.10',
|
||||
'user': 'admin',
|
||||
'port': '22',
|
||||
'identityfile': '~/.ssh/id_ed25519'
|
||||
}
|
||||
}
|
||||
|
||||
Example:
|
||||
>>> hosts = parse_ssh_config()
|
||||
>>> hosts['homelab-1']['hostname']
|
||||
'100.64.1.10'
|
||||
"""
|
||||
if config_path is None:
|
||||
config_path = Path.home() / '.ssh' / 'config'
|
||||
|
||||
if not config_path.exists():
|
||||
logger.warning(f"SSH config not found: {config_path}")
|
||||
return {}
|
||||
|
||||
hosts = {}
|
||||
current_host = None
|
||||
|
||||
try:
|
||||
with open(config_path, 'r') as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
|
||||
# Skip comments and empty lines
|
||||
if not line or line.startswith('#'):
|
||||
continue
|
||||
|
||||
# Host directive
|
||||
if line.lower().startswith('host '):
|
||||
host_alias = line.split(maxsplit=1)[1]
|
||||
# Skip wildcards
|
||||
if '*' not in host_alias and '?' not in host_alias:
|
||||
current_host = host_alias
|
||||
hosts[current_host] = {}
|
||||
|
||||
# Configuration directives
|
||||
elif current_host:
|
||||
parts = line.split(maxsplit=1)
|
||||
if len(parts) == 2:
|
||||
key, value = parts
|
||||
hosts[current_host][key.lower()] = value
|
||||
|
||||
return hosts
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error parsing SSH config: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def parse_sshsync_config(config_path: Optional[Path] = None) -> Dict[str, List[str]]:
|
||||
"""
|
||||
Parse sshsync config file for group definitions.
|
||||
|
||||
Args:
|
||||
config_path: Path to sshsync config (default: ~/.config/sshsync/config.yaml)
|
||||
|
||||
Returns:
|
||||
Dict mapping group names to list of hosts:
|
||||
{
|
||||
'production': ['prod-web-01', 'prod-db-01'],
|
||||
'development': ['dev-laptop', 'dev-desktop']
|
||||
}
|
||||
|
||||
Example:
|
||||
>>> groups = parse_sshsync_config()
|
||||
>>> groups['production']
|
||||
['prod-web-01', 'prod-db-01']
|
||||
"""
|
||||
if config_path is None:
|
||||
config_path = Path.home() / '.config' / 'sshsync' / 'config.yaml'
|
||||
|
||||
if not config_path.exists():
|
||||
logger.warning(f"sshsync config not found: {config_path}")
|
||||
return {}
|
||||
|
||||
try:
|
||||
with open(config_path, 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
return config.get('groups', {})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error parsing sshsync config: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def get_timestamp(iso: bool = True) -> str:
|
||||
"""
|
||||
Get current timestamp.
|
||||
|
||||
Args:
|
||||
iso: If True, return ISO format; otherwise human-readable
|
||||
|
||||
Returns:
|
||||
Timestamp string
|
||||
|
||||
Example:
|
||||
>>> get_timestamp(iso=True)
|
||||
"2025-10-19T19:43:41Z"
|
||||
>>> get_timestamp(iso=False)
|
||||
"2025-10-19 19:43:41"
|
||||
"""
|
||||
now = datetime.now()
|
||||
if iso:
|
||||
return now.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
else:
|
||||
return now.strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
|
||||
def safe_execute(func, *args, default=None, **kwargs) -> Any:
|
||||
"""
|
||||
Execute function with error handling.
|
||||
|
||||
Args:
|
||||
func: Function to execute
|
||||
*args: Positional arguments
|
||||
default: Value to return on error
|
||||
**kwargs: Keyword arguments
|
||||
|
||||
Returns:
|
||||
Function result or default on error
|
||||
|
||||
Example:
|
||||
>>> safe_execute(int, "not_a_number", default=0)
|
||||
0
|
||||
>>> safe_execute(int, "42")
|
||||
42
|
||||
"""
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
except Exception as e:
|
||||
logger.error(f"Error executing {func.__name__}: {e}")
|
||||
return default
|
||||
|
||||
|
||||
def validate_path(path: str, must_exist: bool = True) -> bool:
|
||||
"""
|
||||
Check if path is valid and accessible.
|
||||
|
||||
Args:
|
||||
path: Path to validate
|
||||
must_exist: If True, path must exist
|
||||
|
||||
Returns:
|
||||
True if valid, False otherwise
|
||||
|
||||
Example:
|
||||
>>> validate_path("/tmp")
|
||||
True
|
||||
>>> validate_path("/nonexistent", must_exist=True)
|
||||
False
|
||||
"""
|
||||
p = Path(path).expanduser()
|
||||
|
||||
if must_exist:
|
||||
return p.exists()
|
||||
else:
|
||||
# Check if parent directory exists (for paths that will be created)
|
||||
return p.parent.exists()
|
||||
|
||||
|
||||
def parse_disk_usage(df_output: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Parse 'df' command output.
|
||||
|
||||
Args:
|
||||
df_output: Output from 'df -h' command
|
||||
|
||||
Returns:
|
||||
Dict with disk usage info:
|
||||
{
|
||||
'filesystem': '/dev/sda1',
|
||||
'size': '100G',
|
||||
'used': '45G',
|
||||
'available': '50G',
|
||||
'use_pct': 45,
|
||||
'mount': '/'
|
||||
}
|
||||
|
||||
Example:
|
||||
>>> output = "Filesystem Size Used Avail Use% Mounted on\\n/dev/sda1 100G 45G 50G 45% /"
|
||||
>>> parse_disk_usage(output)
|
||||
{'filesystem': '/dev/sda1', 'size': '100G', ...}
|
||||
"""
|
||||
lines = df_output.strip().split('\n')
|
||||
if len(lines) < 2:
|
||||
return {}
|
||||
|
||||
# Parse last line (actual data, not header)
|
||||
data_line = lines[-1]
|
||||
parts = data_line.split()
|
||||
|
||||
if len(parts) < 6:
|
||||
return {}
|
||||
|
||||
try:
|
||||
return {
|
||||
'filesystem': parts[0],
|
||||
'size': parts[1],
|
||||
'used': parts[2],
|
||||
'available': parts[3],
|
||||
'use_pct': int(parts[4].rstrip('%')),
|
||||
'mount': parts[5]
|
||||
}
|
||||
except (ValueError, IndexError) as e:
|
||||
logger.error(f"Error parsing disk usage: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def parse_memory_usage(free_output: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Parse 'free' command output (Linux).
|
||||
|
||||
Args:
|
||||
free_output: Output from 'free -m' command
|
||||
|
||||
Returns:
|
||||
Dict with memory info:
|
||||
{
|
||||
'total': 16384, # MB
|
||||
'used': 8192,
|
||||
'free': 8192,
|
||||
'use_pct': 50.0
|
||||
}
|
||||
|
||||
Example:
|
||||
>>> output = "Mem: 16384 8192 8192 0 0 0"
|
||||
>>> parse_memory_usage(output)
|
||||
{'total': 16384, 'used': 8192, ...}
|
||||
"""
|
||||
lines = free_output.strip().split('\n')
|
||||
|
||||
for line in lines:
|
||||
if line.startswith('Mem:'):
|
||||
parts = line.split()
|
||||
if len(parts) >= 3:
|
||||
try:
|
||||
total = int(parts[1])
|
||||
used = int(parts[2])
|
||||
free = int(parts[3]) if len(parts) > 3 else (total - used)
|
||||
|
||||
return {
|
||||
'total': total,
|
||||
'used': used,
|
||||
'free': free,
|
||||
'use_pct': (used / total * 100) if total > 0 else 0
|
||||
}
|
||||
except (ValueError, IndexError) as e:
|
||||
logger.error(f"Error parsing memory usage: {e}")
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
def parse_cpu_load(uptime_output: str) -> Dict[str, float]:
|
||||
"""
|
||||
Parse 'uptime' command output for load averages.
|
||||
|
||||
Args:
|
||||
uptime_output: Output from 'uptime' command
|
||||
|
||||
Returns:
|
||||
Dict with load averages:
|
||||
{
|
||||
'load_1min': 0.45,
|
||||
'load_5min': 0.38,
|
||||
'load_15min': 0.32
|
||||
}
|
||||
|
||||
Example:
|
||||
>>> output = "19:43:41 up 5 days, 2:15, 3 users, load average: 0.45, 0.38, 0.32"
|
||||
>>> parse_cpu_load(output)
|
||||
{'load_1min': 0.45, 'load_5min': 0.38, 'load_15min': 0.32}
|
||||
"""
|
||||
# Find "load average:" part
|
||||
match = re.search(r'load average:\s+([\d.]+),\s+([\d.]+),\s+([\d.]+)', uptime_output)
|
||||
|
||||
if match:
|
||||
try:
|
||||
return {
|
||||
'load_1min': float(match.group(1)),
|
||||
'load_5min': float(match.group(2)),
|
||||
'load_15min': float(match.group(3))
|
||||
}
|
||||
except ValueError as e:
|
||||
logger.error(f"Error parsing CPU load: {e}")
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
def format_host_status(host: str, online: bool, groups: List[str],
|
||||
latency: Optional[int] = None,
|
||||
tailscale_connected: bool = False) -> str:
|
||||
"""
|
||||
Format host status as display string.
|
||||
|
||||
Args:
|
||||
host: Host name
|
||||
online: Whether host is online
|
||||
groups: List of groups host belongs to
|
||||
latency: Latency in ms (optional)
|
||||
tailscale_connected: Tailscale connection status
|
||||
|
||||
Returns:
|
||||
Formatted status string
|
||||
|
||||
Example:
|
||||
>>> format_host_status("web-01", True, ["production", "web"], 25, True)
|
||||
"🟢 web-01 (production, web) - Online - Tailscale: Connected | Latency: 25ms"
|
||||
"""
|
||||
icon = "🟢" if online else "🔴"
|
||||
status = "Online" if online else "Offline"
|
||||
group_str = ", ".join(groups) if groups else "no group"
|
||||
|
||||
parts = [f"{icon} {host} ({group_str}) - {status}"]
|
||||
|
||||
if tailscale_connected:
|
||||
parts.append("Tailscale: Connected")
|
||||
|
||||
if latency is not None and online:
|
||||
parts.append(f"Latency: {latency}ms")
|
||||
|
||||
return " - ".join(parts)
|
||||
|
||||
|
||||
def calculate_load_score(cpu_pct: float, mem_pct: float, disk_pct: float) -> float:
|
||||
"""
|
||||
Calculate composite load score for a machine.
|
||||
|
||||
Args:
|
||||
cpu_pct: CPU usage percentage (0-100)
|
||||
mem_pct: Memory usage percentage (0-100)
|
||||
disk_pct: Disk usage percentage (0-100)
|
||||
|
||||
Returns:
|
||||
Load score (0-1, lower is better)
|
||||
|
||||
Formula:
|
||||
score = (cpu * 0.4) + (mem * 0.3) + (disk * 0.3)
|
||||
|
||||
Example:
|
||||
>>> calculate_load_score(45, 60, 40)
|
||||
0.48 # (0.45*0.4 + 0.60*0.3 + 0.40*0.3)
|
||||
"""
|
||||
return (cpu_pct * 0.4 + mem_pct * 0.3 + disk_pct * 0.3) / 100
|
||||
|
||||
|
||||
def classify_load_status(score: float) -> str:
|
||||
"""
|
||||
Classify load score into status category.
|
||||
|
||||
Args:
|
||||
score: Load score (0-1)
|
||||
|
||||
Returns:
|
||||
Status string: "low", "moderate", or "high"
|
||||
|
||||
Example:
|
||||
>>> classify_load_status(0.28)
|
||||
"low"
|
||||
>>> classify_load_status(0.55)
|
||||
"moderate"
|
||||
>>> classify_load_status(0.82)
|
||||
"high"
|
||||
"""
|
||||
if score < 0.4:
|
||||
return "low"
|
||||
elif score < 0.7:
|
||||
return "moderate"
|
||||
else:
|
||||
return "high"
|
||||
|
||||
|
||||
def classify_latency(latency_ms: int) -> Tuple[str, str]:
|
||||
"""
|
||||
Classify network latency.
|
||||
|
||||
Args:
|
||||
latency_ms: Latency in milliseconds
|
||||
|
||||
Returns:
|
||||
Tuple of (status, description)
|
||||
|
||||
Example:
|
||||
>>> classify_latency(25)
|
||||
("excellent", "Ideal for interactive tasks")
|
||||
>>> classify_latency(150)
|
||||
("fair", "May impact interactive workflows")
|
||||
"""
|
||||
if latency_ms < 50:
|
||||
return ("excellent", "Ideal for interactive tasks")
|
||||
elif latency_ms < 100:
|
||||
return ("good", "Suitable for most operations")
|
||||
elif latency_ms < 200:
|
||||
return ("fair", "May impact interactive workflows")
|
||||
else:
|
||||
return ("poor", "Investigate network issues")
|
||||
|
||||
|
||||
def get_hosts_from_groups(group: str, groups_config: Dict[str, List[str]]) -> List[str]:
|
||||
"""
|
||||
Get list of hosts in a group.
|
||||
|
||||
Args:
|
||||
group: Group name
|
||||
groups_config: Groups configuration dict
|
||||
|
||||
Returns:
|
||||
List of host names in group
|
||||
|
||||
Example:
|
||||
>>> groups = {'production': ['web-01', 'db-01']}
|
||||
>>> get_hosts_from_groups('production', groups)
|
||||
['web-01', 'db-01']
|
||||
"""
|
||||
return groups_config.get(group, [])
|
||||
|
||||
|
||||
def get_groups_for_host(host: str, groups_config: Dict[str, List[str]]) -> List[str]:
|
||||
"""
|
||||
Get list of groups a host belongs to.
|
||||
|
||||
Args:
|
||||
host: Host name
|
||||
groups_config: Groups configuration dict
|
||||
|
||||
Returns:
|
||||
List of group names
|
||||
|
||||
Example:
|
||||
>>> groups = {'production': ['web-01'], 'web': ['web-01', 'web-02']}
|
||||
>>> get_groups_for_host('web-01', groups)
|
||||
['production', 'web']
|
||||
"""
|
||||
return [group for group, hosts in groups_config.items() if host in hosts]
|
||||
|
||||
|
||||
def run_command(command: str, timeout: int = 10) -> Tuple[bool, str, str]:
|
||||
"""
|
||||
Run shell command with timeout.
|
||||
|
||||
Args:
|
||||
command: Command to execute
|
||||
timeout: Timeout in seconds
|
||||
|
||||
Returns:
|
||||
Tuple of (success, stdout, stderr)
|
||||
|
||||
Example:
|
||||
>>> success, stdout, stderr = run_command("echo hello")
|
||||
>>> success
|
||||
True
|
||||
>>> stdout.strip()
|
||||
"hello"
|
||||
"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
command,
|
||||
shell=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
return (
|
||||
result.returncode == 0,
|
||||
result.stdout,
|
||||
result.stderr
|
||||
)
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
return (False, "", f"Command timed out after {timeout}s")
|
||||
except Exception as e:
|
||||
return (False, "", str(e))
|
||||
|
||||
|
||||
def main():
|
||||
"""Test helper functions."""
|
||||
print("Testing helper functions...\n")
|
||||
|
||||
# Test formatting
|
||||
print("1. Format bytes:")
|
||||
print(f" 12582912 bytes = {format_bytes(12582912)}")
|
||||
print(f" 1610612736 bytes = {format_bytes(1610612736)}")
|
||||
|
||||
print("\n2. Format duration:")
|
||||
print(f" 135 seconds = {format_duration(135)}")
|
||||
print(f" 5430 seconds = {format_duration(5430)}")
|
||||
|
||||
print("\n3. Format percentage:")
|
||||
print(f" 45.567 = {format_percentage(45.567)}")
|
||||
|
||||
print("\n4. Calculate load score:")
|
||||
score = calculate_load_score(45, 60, 40)
|
||||
print(f" CPU 45%, Mem 60%, Disk 40% = {score:.2f}")
|
||||
print(f" Status: {classify_load_status(score)}")
|
||||
|
||||
print("\n5. Classify latency:")
|
||||
latencies = [25, 75, 150, 250]
|
||||
for lat in latencies:
|
||||
status, desc = classify_latency(lat)
|
||||
print(f" {lat}ms: {status} - {desc}")
|
||||
|
||||
print("\n6. Parse SSH config:")
|
||||
ssh_hosts = parse_ssh_config()
|
||||
print(f" Found {len(ssh_hosts)} hosts")
|
||||
|
||||
print("\n7. Parse sshsync config:")
|
||||
groups = parse_sshsync_config()
|
||||
print(f" Found {len(groups)} groups")
|
||||
for group, hosts in groups.items():
|
||||
print(f" - {group}: {len(hosts)} hosts")
|
||||
|
||||
print("\n✅ All helpers tested successfully")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user