|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import getpass |
| 4 | +import os |
| 5 | +import shlex |
| 6 | +import subprocess as sp |
| 7 | +import tempfile |
| 8 | +import sys |
| 9 | +from pathlib import Path |
| 10 | + |
| 11 | +from ruamel.yaml import YAML |
| 12 | + |
| 13 | +from .inventory import ROOT, ansible_inventory, save |
| 14 | +from .models import Host, Inventory |
| 15 | + |
| 16 | + |
| 17 | +def check_admin(host: Host, *, report_error: bool = False) -> bool: |
| 18 | + command = [ |
| 19 | + 'ssh', '-o', 'BatchMode=yes', '-o', 'IdentitiesOnly=yes', |
| 20 | + '-o', 'PreferredAuthentications=publickey', |
| 21 | + '-o', 'ControlMaster=no', '-o', 'ControlPath=none', '-o', 'ConnectTimeout=10', |
| 22 | + '-p', str(host.admin.port), '-i', str(Path(host.admin.private_key_file).expanduser()), |
| 23 | + *shlex.split(os.environ.get('ANSIBLE_SSH_ARGS', '')), |
| 24 | + *shlex.split(os.environ.get('ANSIBLE_SSH_COMMON_ARGS', '')), |
| 25 | + *shlex.split(os.environ.get('ANSIBLE_SSH_EXTRA_ARGS', '')), |
| 26 | + '-l', host.admin.user, '--', host.address, 'sudo -n true', |
| 27 | + ] |
| 28 | + try: |
| 29 | + result = sp.run(command, capture_output=True, text=True, timeout=30) |
| 30 | + except sp.TimeoutExpired: |
| 31 | + if report_error: |
| 32 | + print('Administrative SSH probe timed out after 30 seconds', file=sys.stderr) |
| 33 | + return False |
| 34 | + if result.returncode and report_error: |
| 35 | + print(f'Administrative SSH probe failed (exit {result.returncode}):', file=sys.stderr) |
| 36 | + print(result.stderr.strip() or result.stdout.strip() or 'SSH returned no diagnostic output', file=sys.stderr) |
| 37 | + return result.returncode == 0 |
| 38 | + |
| 39 | + |
| 40 | +def run_phase(inventory: Inventory, name: str, phase: str, password: str | None = None) -> int: |
| 41 | + host = inventory.hosts[name] |
| 42 | + data = ansible_inventory(inventory) |
| 43 | + variables = data['all']['hosts'][name] |
| 44 | + data['all']['hosts'] = {name: variables} |
| 45 | + variables['ansible_user'] = host.admin.user if phase == 'policy' else host.admin.bootstrap_user |
| 46 | + variables['ansible_ssh_private_key_file'] = ( |
| 47 | + host.admin.private_key_file if phase == 'policy' |
| 48 | + else host.admin.bootstrap_private_key_file or host.admin.private_key_file |
| 49 | + ) |
| 50 | + if password is not None: |
| 51 | + variables['ansible_password'] = password |
| 52 | + variables['ansible_become_password'] = password |
| 53 | + variables['ansible_ssh_common_args'] = '-o PubkeyAuthentication=no -o PreferredAuthentications=password,keyboard-interactive' |
| 54 | + variables['megaproxy_bootstrap_phase'] = phase |
| 55 | + # Bootstrap passwords exist only in a private temporary directory for this subprocess. |
| 56 | + with tempfile.TemporaryDirectory(prefix='megaproxy-bootstrap-') as directory: |
| 57 | + path = Path(directory) / 'inventory.yml' |
| 58 | + with open(path, 'w', encoding='utf-8', opener=lambda name, flags: os.open(name, flags, 0o600)) as stream: |
| 59 | + YAML().dump(data, stream) |
| 60 | + return sp.run([ |
| 61 | + 'ansible-playbook', '-i', str(path), str(ROOT / 'playbooks/bootstrap.yml'), |
| 62 | + ], cwd=ROOT).returncode |
| 63 | + |
| 64 | + |
| 65 | +def bootstrap(path: Path, inventory: Inventory, limit: str | None = None) -> int: |
| 66 | + names = list(inventory.hosts) if not limit else [name.strip() for name in limit.split(',')] |
| 67 | + if not names or any(name not in inventory.hosts for name in names): |
| 68 | + raise ValueError('Bootstrap requires exact inventory host names in --limit') |
| 69 | + for name in dict.fromkeys(names): |
| 70 | + host = inventory.hosts[name] |
| 71 | + if not host.admin.bootstrap_user: |
| 72 | + continue |
| 73 | + print(f'Bootstrapping administrative access: {name}', flush=True) |
| 74 | + if not check_admin(host): |
| 75 | + password = None |
| 76 | + if host.admin.bootstrap_auth == 'password': |
| 77 | + password = getpass.getpass(f'Initial SSH password for {host.admin.bootstrap_user}@{host.address}: ') |
| 78 | + if not password: |
| 79 | + raise ValueError('Initial SSH password must not be empty') |
| 80 | + code = run_phase(inventory, name, 'account', password) |
| 81 | + password = None |
| 82 | + if code: |
| 83 | + return code |
| 84 | + if not check_admin(host, report_error=True): |
| 85 | + raise ValueError(f'Administrative key login or sudo failed for {name}; SSH policy was not changed') |
| 86 | + code = run_phase(inventory, name, 'policy') |
| 87 | + if code: |
| 88 | + return code |
| 89 | + if not check_admin(host, report_error=True): |
| 90 | + raise ValueError(f'Administrative access verification failed after SSH reload for {name}; bootstrap remains pending') |
| 91 | + host.admin.bootstrap_user = None |
| 92 | + host.admin.bootstrap_private_key_file = None |
| 93 | + host.admin.bootstrap_auth = 'key' |
| 94 | + save(path, inventory) |
| 95 | + print(f'Administrative connection activated for: {name}', flush=True) |
| 96 | + return 0 |
0 commit comments