alexveley
399 lines
13 KiB
Bash
399 lines
13 KiB
Bash
#!/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'''
|
|
$ProgressPreference = 'SilentlyContinue'
|
|
$ErrorActionPreference = 'SilentlyContinue'
|
|
$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 clean_exec_text(text: str) -> str:
|
|
if not text:
|
|
return ''
|
|
text = re.sub(r'#< CLIXML.*', '', text, flags=re.DOTALL)
|
|
return text.strip()
|
|
|
|
|
|
def extract_id_from_exec(out: str, err: str) -> str | None:
|
|
combined = '\n'.join(filter(None, (clean_exec_text(out), clean_exec_text(err))))
|
|
for line in combined.splitlines():
|
|
line = line.strip()
|
|
if re.fullmatch(r'\d{6,12}', line):
|
|
return line
|
|
match = re.search(r'\d{6,12}', combined)
|
|
return match.group(0) if match else None
|
|
|
|
|
|
def parse_exec_result(st: dict) -> tuple[str | None, str | None]:
|
|
try:
|
|
exitcode = int(st.get('exitcode', 1))
|
|
out, err = decode_exec_payload(st)
|
|
rid = extract_id_from_exec(out, err)
|
|
if rid:
|
|
return rid, None
|
|
if exitcode != 0:
|
|
brief = clean_exec_text(err) or clean_exec_text(out) or f'guest exit {exitcode}'
|
|
if len(brief) > 120:
|
|
brief = brief[:120] + '...'
|
|
return None, brief
|
|
return None, f'no id in exec output: {out!r}'
|
|
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', '-NonInteractive', '-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
|