Codesys-MCP/python_scripting/deploy_winrm_agent.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import base64
import pathlib
import sys
from typing import Iterable

import winrm


def run_ps(session: winrm.Session, script: str) -> str:
    result = session.run_ps(script)
    stdout = result.std_out.decode("utf-8", errors="ignore").strip()
    stderr = result.std_err.decode("utf-8", errors="ignore").strip()
    if result.status_code != 0:
        raise RuntimeError(f"PowerShell failed ({result.status_code}): {stderr or stdout}")
    return stdout


def upload_file(session: winrm.Session, local_path: pathlib.Path, remote_path: str) -> None:
    data_b64 = base64.b64encode(local_path.read_bytes()).decode("ascii")
    temp_b64_path = remote_path + ".b64tmp"
    # WinRM command payloads can be short; use conservative chunking.
    chunk_size = 500

    run_ps(
        session,
        (
            f"$p='{remote_path}';"
            "$d=Split-Path -Parent $p;"
            "New-Item -ItemType Directory -Force -Path $d | Out-Null;"
            f"[System.IO.File]::WriteAllText('{temp_b64_path}','');"
            "Write-Output 'ready'"
        ),
    )

    for i in range(0, len(data_b64), chunk_size):
        chunk = data_b64[i : i + chunk_size]
        run_ps(
            session,
            f"Add-Content -Path '{temp_b64_path}' -Value '{chunk}' -NoNewline",
        )

    ps_finalize = (
        f"$b64=[System.IO.File]::ReadAllText('{temp_b64_path}');"
        f"[System.IO.File]::WriteAllBytes('{remote_path}',[Convert]::FromBase64String($b64));"
        f"Remove-Item -Path '{temp_b64_path}' -Force;"
        f"Write-Output ('uploaded:' + '{remote_path}')"
    )
    print(run_ps(session, ps_finalize))


def iter_agent_files(root: pathlib.Path) -> Iterable[pathlib.Path]:
    for p in sorted(root.rglob("*")):
        if p.is_file():
            yield p


def main() -> int:
    parser = argparse.ArgumentParser(description="Deploy Windows CODESYS agent via WinRM.")
    parser.add_argument("--host", required=True)
    parser.add_argument("--username", required=True)
    parser.add_argument("--password", required=True)
    parser.add_argument("--token", required=True)
    parser.add_argument("--install-path", default=r"C:\codesys-agent")
    parser.add_argument("--codesys-exe-path", required=True)
    parser.add_argument("--codesys-profile", required=True)
    args = parser.parse_args()

    session = winrm.Session(
        target=f"http://{args.host}:5985/wsman",
        auth=(args.username, args.password),
        transport="ntlm",
        server_cert_validation="ignore",
    )

    print("Connected, checking identity...")
    whoami = run_ps(session, "whoami")
    print(f"Remote identity: {whoami}")

    local_root = pathlib.Path(__file__).resolve().parent / "windows_codesys_agent"
    if not local_root.is_dir():
        raise RuntimeError(f"Local folder not found: {local_root}")

    remote_root = args.install_path.rstrip("\\") + r"\windows_codesys_agent"
    for f in iter_agent_files(local_root):
        rel = f.relative_to(local_root).as_posix().replace("/", "\\")
        remote_path = remote_root + "\\" + rel
        upload_file(session, f, remote_path)

    scripts_dir = remote_root + r"\scripts"
    runtime_dir = args.install_path.rstrip("\\") + r"\runtime"
    run_script_path = remote_root + r"\run-agent.ps1"
    run_script = (
        f"$env:CODESYS_AGENT_TOKEN = '{args.token}'\n"
        f"$env:CODESYS_AGENT_SCRIPTS_DIR = '{scripts_dir}'\n"
        f"$env:CODESYS_AGENT_RUNTIME_DIR = '{runtime_dir}'\n"
        "$env:CODESYS_AGENT_ALLOW_DOWNLOAD = 'false'\n"
        "$env:CODESYS_AGENT_TIMEOUT_SEC = '180'\n"
        f"$env:CODESYS_EXE_PATH = '{args.codesys_exe_path}'\n"
        f"$env:CODESYS_PROFILE = '{args.codesys_profile}'\n"
        "$env:CODESYS_NO_UI = 'true'\n"
        "$env:CODESYS_SCRIPT_TIMEOUT_SEC = '600'\n"
        "Set-Location -Path $PSScriptRoot\n"
        ".\\.venv\\Scripts\\python.exe -m uvicorn agent:app --host 0.0.0.0 --port 8787\n"
    )
    run_b64 = base64.b64encode(run_script.encode("utf-8")).decode("ascii")
    run_ps(
        session,
        (
            f"$p='{run_script_path}';"
            f"[System.IO.File]::WriteAllBytes($p,[Convert]::FromBase64String('{run_b64}'));"
            "Write-Output ('uploaded:' + $p)"
        ),
    )

    setup_ps = (
        f"Set-Location -Path '{remote_root}';"
        "if (Get-Command py -ErrorAction SilentlyContinue) { py -3 -m venv .venv } "
        "elseif (Get-Command python -ErrorAction SilentlyContinue) { python -m venv .venv } "
        "else { throw 'Python is not installed on the Windows VM.' };"
        ".\\.venv\\Scripts\\python.exe -m pip install -r requirements.txt;"
        "Get-Process -Name python -ErrorAction SilentlyContinue | "
        "Where-Object { $_.Path -like '*windows_codesys_agent*' } | Stop-Process -Force -ErrorAction SilentlyContinue;"
        "Start-Process -FilePath powershell.exe "
        f"-ArgumentList '-ExecutionPolicy Bypass -File \"{run_script_path}\"' "
        "-WindowStyle Hidden;"
        "Write-Output 'agent-started'"
    )
    print(run_ps(session, setup_ps))
    print(f"Deployment complete. Agent expected at: http://{args.host}:8787/health")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        raise