#!/usr/bin/env python3
"""Roda um comando no guest atraves do console serial do libvirt.

O console serial e o unico canal garantido nas imagens cloud minimas: nao
depende de sshd, de cloud-init nem de rede. Este script sincroniza com o
prompt, autentica (root/rfstest) e executa o comando, lendo a saida do pty
em blocos pequenos para evitar a sobreposicao de eco que ocorre em comandos
longos.

Uso:
    guestcmd.py <dominio> '<comando>'
    guestcmd.py <dominio> -f <arquivo>       # executa um script local
"""
import base64
import os
import re
import select
import subprocess
import sys
import time

USER = "root"
PASSWORD = "rfstest"


def console_pty(domain):
    xml = subprocess.run(["virsh", "dumpxml", domain],
                         capture_output=True, text=True).stdout
    m = (re.search(r"<console[^>]*>\s*<source path='([^']+)'", xml)
         or re.search(r"<serial[^>]*>\s*<source path='([^']+)'", xml))
    if not m:
        sys.exit(f"console serial nao encontrada para {domain}")
    return m.group(1)


class Console:
    def __init__(self, path):
        self.fd = os.open(path, os.O_RDWR | os.O_NOCTTY)

    def send(self, s):
        os.write(self.fd, s.encode())

    def drain(self, seconds):
        buf = b""
        end = time.time() + seconds
        while time.time() < end:
            r, _, _ = select.select([self.fd], [], [], 0.3)
            if not r:
                continue
            try:
                chunk = os.read(self.fd, 65536)
            except OSError:
                break
            if not chunk:
                break
            buf += chunk
        return buf

    def login(self):
        self.send("\n")
        time.sleep(1.0)
        snap = self.drain(2.5).decode(errors="replace").lower()
        if "login:" not in snap:
            return          # ja autenticado
        self.send(USER + "\n")
        self.drain(3.0)
        self.send(PASSWORD + "\n")
        time.sleep(2.0)
        self.drain(2.0)

    def run(self, command, timeout=120.0):
        """Executa e devolve a saida, sincronizando por marcador no pty."""
        out_path = "/tmp/.guestcmd.out"
        self.send(f"rm -f {out_path}; ( {command} ) > {out_path} 2>&1; "
                  f"echo __GC_DONE__\n")
        deadline = time.time() + timeout
        while time.time() < deadline:
            if b"__GC_DONE__" in self.drain(2.0):
                break
        else:
            return "(timeout)\n"

        # tamanho em bytes
        self.send(f"wc -c < {out_path}\n")
        time.sleep(1.2)
        raw = self.drain(2.0).decode(errors="replace")
        nums = re.findall(r"\d+", raw)
        size = int(nums[-1]) if nums else 0

        # ler em blocos base64 (evita eco/sobreposicao)
        data = b""
        while len(data) < size:
            self.send(f"dd if={out_path} bs=1 skip={len(data)} count=150 "
                      f"2>/dev/null | base64 -w0; echo\n")
            time.sleep(0.8)
            enc = self.drain(2.0).decode(errors="replace")
            token = "".join(re.findall(r"[A-Za-z0-9+/=]{4,}", enc))
            if not token:
                break
            try:
                data += base64.b64decode(token + "=" * (-len(token) % 4))
            except Exception:
                break
        return data.decode(errors="replace")


def main():
    if len(sys.argv) < 3:
        sys.exit(__doc__)
    domain = sys.argv[1]
    console = Console(console_pty(domain))
    console.login()

    if sys.argv[2] == "-f":
        with open(sys.argv[3]) as fh:
            script = fh.read().replace("'", "'\\''")
        command = f"sh -c '{script}'"
    else:
        command = sys.argv[2]

    sys.stdout.write(console.run(command))
    sys.stdout.write("\n")


if __name__ == "__main__":
    main()