#!/bin/bash
# Poll RustDesk IDs via QEMU guest agent; update Proxmox VM notes + registry.json.
# Install to /usr/local/sbin on pve. Cron every 5m. No guest reporter or share reports needed.
# Usage: blair-sync-rustdesk-via-ga.sh [--dry-run]
set -euo pipefail

export BLAIR_SHARE="${BLAIR_SHARE:-/tank/blair-share}"
[[ "${1:-}" == "--dry-run" ]] && export BLAIR_DRY_RUN=1

if ! command -v python3 >/dev/null 2>&1; then
  echo 'python3 required on pve' >&2
  exit 1
fi

exec python3 - <<'PY'
import base64
import binascii
import json
import os
import re
import subprocess
import time
from datetime import date

SHARE = os.environ.get('BLAIR_SHARE', '/tank/blair-share')
REGISTRY = os.path.join(SHARE, 'fleet', 'registry.json')
NODE = os.environ.get('BLAIR_PVE_NODE', 'pve')
DRY_RUN = os.environ.get('BLAIR_DRY_RUN', '') in ('1', 'true', 'yes')
VERBOSE = os.environ.get('BLAIR_VERBOSE', '') in ('1', 'true', 'yes')
VMIDS_ENV = os.environ.get('BLAIR_VMIDS', '').strip()

GUEST_READ_PATHS = [
    r'C:\ProgramData\Blair\rustdesk-report.json',
    r'C:\Windows\ServiceProfiles\LocalService\AppData\Roaming\RustDesk\config\RustDesk2.toml',
    r'C:\Windows\ServiceProfiles\LocalService\AppData\Roaming\RustDesk\config\RustDesk.toml',
    r'C:\Windows\System32\config\systemprofile\AppData\Roaming\RustDesk\config\RustDesk2.toml',
    r'C:\ProgramData\RustDesk\config\RustDesk2.toml',
]

PS_READ_ID = r'''
$paths = @(
  'C:\ProgramData\Blair\rustdesk-report.json',
  'C:\Windows\ServiceProfiles\LocalService\AppData\Roaming\RustDesk\config\RustDesk2.toml',
  'C:\Windows\ServiceProfiles\LocalService\AppData\Roaming\RustDesk\config\RustDesk.toml',
  'C:\Windows\System32\config\systemprofile\AppData\Roaming\RustDesk\config\RustDesk2.toml',
  'C:\ProgramData\RustDesk\config\RustDesk2.toml'
)
foreach ($p in $paths) {
  if (-not (Test-Path -LiteralPath $p)) { continue }
  $t = Get-Content -Raw -LiteralPath $p
  if ($p -like '*.json') {
    try {
      $j = $t | ConvertFrom-Json
      if ($j.rustdesk_id -match '^\d{6,12}$') { Write-Output $j.rustdesk_id; exit 0 }
    } catch {}
  }
  if ($t -match '(?m)^\s*id\s*=\s*[''"]?(\d{6,12})') { Write-Output $Matches[1]; exit 0 }
}
foreach ($rp in @(
  'HKLM:\SOFTWARE\RustDesk\Host',
  'HKLM:\SOFTWARE\RustDesk',
  'HKLM:\SOFTWARE\WOW6432Node\RustDesk\Host',
  'HKLM:\SOFTWARE\WOW6432Node\RustDesk'
)) {
  if (-not (Test-Path $rp)) { continue }
  foreach ($n in @('id', 'Id', 'rustdesk_id', 'peer_id')) {
    $v = (Get-ItemProperty -Path $rp -Name $n -ErrorAction SilentlyContinue).$n
    if ($v -match '^\d{6,12}$') { Write-Output $v; exit 0 }
  }
}
exit 1
'''


def log(msg: str) -> None:
    print(msg, flush=True)


def update_desc(old: str, new_id: str) -> str:
    if re.search(r'RustDesk\s+\d+', old):
        return re.sub(r'RustDesk\s+\d+', f'RustDesk {new_id}', old)
    if not old:
        return f'RustDesk {new_id}.'
    return f"{old.rstrip('.')}. RustDesk {new_id}."


def load_registry() -> dict:
    with open(REGISTRY, encoding='utf-8') as fh:
        return json.load(fh)


def fleet_vmids(registry: dict) -> list[int]:
    if VMIDS_ENV:
        return [int(x) for x in VMIDS_ENV.split() if x.isdigit()]
    return [int(s['vmid']) for s in registry.get('windows_gpu_seats', []) if s.get('vmid')]


def registry_old_id(registry: dict, vmid: int) -> str:
    for seat in registry.get('windows_gpu_seats', []):
        if seat.get('vmid') == vmid:
            return str(seat.get('rustdesk_id', '') or '')
    return ''


def qm_running(vmid: int) -> bool:
    proc = subprocess.run(['qm', 'status', str(vmid)], capture_output=True, text=True, check=False)
    return 'status: running' in proc.stdout


def guest_ping(vmid: int) -> bool:
    proc = subprocess.run(
        ['qm', 'guest', 'cmd', str(vmid), 'ping'],
        capture_output=True, text=True, check=False,
    )
    return proc.returncode == 0


def decode_ga_payload(raw) -> str:
    if raw is None:
        return ''
    if isinstance(raw, bytes):
        return raw.decode('utf-8-sig', errors='replace').strip()
    text = str(raw).strip()
    if not text:
        return ''
    if text[0] in '{[#' or 'rustdesk_id' in text or re.search(r'(?m)^\s*id\s*=', text):
        return text.lstrip('\ufeff')
    if re.fullmatch(r'[\d\r\n\t ]+', text):
        return text.strip()
    blob = ''.join(text.split())
    if not blob:
        return text
    for candidate in (blob, blob + ('=' * ((4 - len(blob) % 4) % 4))):
        try:
            decoded = base64.b64decode(candidate, validate=False)
            return decoded.decode('utf-8-sig', errors='replace').strip()
        except (ValueError, binascii.Error):
            continue
    return text


def decode_exec_payload(st: dict) -> tuple[str, str]:
    out = decode_ga_payload(st.get('out-data'))
    err = decode_ga_payload(st.get('err-data'))
    return out, err


def parse_exec_result(st: dict) -> tuple[str | None, str | None]:
    try:
        exitcode = int(st.get('exitcode', 1))
        out, err = decode_exec_payload(st)
        if exitcode != 0:
            return None, err or out or f'guest exit {exitcode}'
        match = re.search(r'\d{6,12}', out)
        if not match:
            return None, f'no id in exec output: {out!r}'
        return match.group(0), None
    except Exception as exc:  # noqa: BLE001
        return None, f'exec decode error: {exc}'


def guest_exec_text(vmid: int, command: list[str], timeout: int = 90) -> tuple[str | None, str | None]:
    proc = subprocess.run(
        ['qm', 'guest', 'exec', str(vmid), '--'] + command,
        capture_output=True, text=True, check=False,
    )
    if not proc.stdout.strip():
        return None, proc.stderr.strip() or f'qm guest exec failed rc={proc.returncode}'
    try:
        data = json.loads(proc.stdout)
    except json.JSONDecodeError:
        return None, f'exec json parse error: {proc.stdout[:160]!r}'

    if data.get('exited') == 1:
        return parse_exec_result(data)

    pid = data.get('pid')
    if not pid:
        return None, f'exec returned no pid: {data!r}'

    deadline = time.time() + timeout
    while time.time() < deadline:
        st_proc = subprocess.run(
            ['qm', 'guest', 'exec-status', str(vmid), str(pid)],
            capture_output=True, text=True, check=False,
        )
        try:
            st = json.loads(st_proc.stdout)
        except json.JSONDecodeError:
            time.sleep(2)
            continue
        if st.get('exited') != 1:
            time.sleep(2)
            continue
        return parse_exec_result(st)

    return None, 'guest exec timeout'


def decode_agent_file_content(raw) -> str | None:
    text = decode_ga_payload(raw)
    return text or None


def parse_id_from_content(text: str, path: str) -> str | None:
    if path.lower().endswith('.json'):
        try:
            rid = str(json.loads(text).get('rustdesk_id', '') or '').strip()
            if re.fullmatch(r'\d{6,12}', rid):
                return rid
        except json.JSONDecodeError:
            pass
        match = re.search(r'"rustdesk_id"\s*:\s*"(\d{6,12})"', text)
        if match:
            return match.group(1)
    for pattern in (
        r'(?m)^\s*id\s*=\s*["\']?(\d{6,12})',
        r'(?m)rustdesk_id\s*=\s*["\']?(\d{6,12})',
    ):
        match = re.search(pattern, text)
        if match:
            return match.group(1)
    return None


def guest_file_read(vmid: int, guest_path: str) -> tuple[str | None, str | None]:
    for path in {guest_path, guest_path.replace('\\', '/')}:
        try:
            proc = subprocess.run(
                [
                    'pvesh', 'get', f'/nodes/{NODE}/qemu/{vmid}/agent/file-read',
                    '--file', path, '--output-format', 'json',
                ],
                capture_output=True, text=True, check=False,
            )
            if proc.returncode != 0:
                continue
            payload = json.loads(proc.stdout)
            data = payload.get('data', payload)
            if not isinstance(data, dict):
                continue
            raw = data.get('content')
            if raw is None:
                continue
            text = decode_agent_file_content(raw)
            if text:
                return text, None
        except (json.JSONDecodeError, ValueError, UnicodeError):
            continue
    return None, f'file-read failed for {guest_path}'


def powershell_encoded(script: str) -> str:
    return base64.b64encode(script.strip().encode('utf-16-le')).decode('ascii')


def guest_exec_powershell(vmid: int) -> tuple[str | None, str | None]:
    enc = powershell_encoded(PS_READ_ID)
    return guest_exec_text(
        vmid,
        ['powershell.exe', '-NoProfile', '-EncodedCommand', enc],
        timeout=120,
    )


def read_rustdesk_id(vmid: int) -> tuple[str | None, str | None]:
    errors: list[str] = []
    for path in GUEST_READ_PATHS:
        try:
            text, err = guest_file_read(vmid, path)
            if not text:
                if VERBOSE:
                    errors.append(f'{path}: {err}')
                continue
            rid = parse_id_from_content(text, path)
            if rid:
                if VERBOSE:
                    log(f'VM {vmid}: ID {rid} from file-read {path}')
                return rid, None
            errors.append(f'{path}: no id in file')
        except Exception as exc:  # noqa: BLE001 — keep sync loop alive
            errors.append(f'{path}: {exc}')
            continue

    rid, err = guest_exec_powershell(vmid)
    if rid:
        if VERBOSE:
            log(f'VM {vmid}: ID {rid} from guest exec')
        return rid, None
    errors.append(f'exec: {err}')
    return None, '; '.join(errors[-4:]) if errors else 'unknown'


def qm_description(vmid: int) -> str:
    proc = subprocess.run(['qm', 'config', str(vmid)], capture_output=True, text=True, check=False)
    for line in proc.stdout.splitlines():
        if line.startswith('description:'):
            return line.split(':', 1)[1].strip()
    return ''


def desc_rustdesk_id(desc: str) -> str:
    match = re.search(r'RustDesk\s+(\d{6,12})', desc)
    return match.group(1) if match else ''


def save_registry(registry: dict) -> None:
    tmp = REGISTRY + '.tmp'
    with open(tmp, 'w', encoding='utf-8') as fh:
        json.dump(registry, fh, indent=2)
        fh.write('\n')
    os.replace(tmp, REGISTRY)


def main() -> int:
    if not os.path.isfile(REGISTRY):
        log(f'Missing registry: {REGISTRY}')
        return 1

    if DRY_RUN:
        log('DRY RUN - no qm set or registry writes')

    registry = load_registry()
    vmids = fleet_vmids(registry)
    if not vmids:
        log('No fleet VMIDs found')
        return 0

    registry_dirty = False
    for vmid in vmids:
        if not qm_running(vmid):
            log(f'VM {vmid}: skip (not running)')
            continue
        if not guest_ping(vmid):
            log(f'VM {vmid}: skip (guest agent not responding)')
            continue

        new_id, err = read_rustdesk_id(vmid)
        if not new_id:
            hint = ' (try BLAIR_VERBOSE=1)' if not VERBOSE else ''
            log(f'VM {vmid}: skip (could not read RustDesk ID: {err}){hint}')
            continue

        old_registry_id = registry_old_id(registry, vmid)
        current_desc = qm_description(vmid)
        old_desc_id = desc_rustdesk_id(current_desc)

        if new_id == old_registry_id and new_id == old_desc_id:
            log(f'VM {vmid}: OK ({new_id})')
            continue

        newdesc = update_desc(current_desc, new_id)
        log(f'VM {vmid}: {old_registry_id or old_desc_id or "(none)"} -> {new_id}')

        if DRY_RUN:
            preview = newdesc if len(newdesc) <= 120 else newdesc[:120] + '...'
            log(f'  would set description: {preview}')
            continue

        subprocess.run(['qm', 'set', str(vmid), '-description', newdesc], check=True)
        registry['updated'] = date.today().isoformat()
        for seat in registry.get('windows_gpu_seats', []):
            if seat.get('vmid') == vmid:
                seat['rustdesk_id'] = new_id
                registry_dirty = True

    if registry_dirty and not DRY_RUN:
        save_registry(registry)
        log('Updated registry.json')

    return 0


if __name__ == '__main__':
    raise SystemExit(main())
PY