META - Attach Enhancements

#!/usr/bin/env pwsh

using namespace System.IO
using namespace System.IO.Pipes
using namespace System.Management.Automation.Language
using namespace System.Net.Sockets
using namespace System.Text
using namespace System.Threading
using namespace System.Threading.Tasks

#Requires -Version 7.4

[CmdletBinding()]
param ()

$ErrorActionPreference = 'Stop'

$configPath = '~/.ansible/test/debugging'
if (-not (Test-Path -LiteralPath $configPath)) {
    New-Item -ItemType Directory -Path $configPath | Out-Null
}
$configFile = Join-Path -Path $configPath -ChildPath 'pwsh_debugger.json'

$sock = $null
try {
    # Start listening to a socket locally that we will forward to the remote
    # host via SSH.
    $sock = [TcpListener]::new(
        [IPAddress]::IPv6Loopback,
        0)
    $sock.Server.DualMode = $true
    $sock.Start()

    $configInfo = @{
        pid = $PID
        port = $sock.LocalEndpoint.Port
    } | ConvertTo-Json -Compress
    Set-Content -LiteralPath $configFile -Value $configInfo

    Write-Host "Listening for debug request on port $($sock.LocalEndpoint.Port)"

    while ($true) {
        $task = $sock.AcceptTcpClientAsync()
        while (-not $task.AsyncWaitHandle.WaitOne(300)) { }
        $client = $task.GetAwaiter().GetResult()

        Write-Host "Client connected from $($client.Client.RemoteEndPoint)"
        $pipe = $stream = $null
        try {
            $stream = $client.GetStream()

            $reader = [StreamReader]::new($stream, [Encoding]::UTF8, $true, $client.ReceiveBufferSize, $true)
            $task = $reader.ReadLineAsync()
            while (-not $task.AsyncWaitHandle.WaitOne(300)) { }
            $debugInfoRaw = $task.GetAwaiter().GetResult()
            $reader.Dispose()

            Write-Host "Received debug request`n$debugInfoRaw"
            $debugInfo = ConvertFrom-Json -InputObject $debugInfoRaw

            $pipeName = "MyPipe-$(New-Guid)"
            $pipe = [NamedPipeServerStream]::new(
                $pipeName,
                [PipeDirection]::InOut,
                1,
                [PipeTransmissionMode]::Byte,
                [PipeOptions]::Asynchronous)

            $task = $pipe.WaitForConnectionAsync()

            Write-Host "Starting VSCode attach to Pipe $pipeName"
            $attachConfig = @{
                name = $debugInfo.ModuleName
                type = "PowerShell"
                request = "attach"
                customPipeName = $pipeName
                runspaceId = $debugInfo.RunspaceId
                pathMappings = @($debugInfo.PathMappings)
            }
            $attachTask = Start-NewDebugSession -Request Attach -Configuration $attachConfig

            Write-Host "Waiting for VSCode to attach to Pipe"
            while (-not $task.AsyncWaitHandle.WaitOne(300)) {}
            $null = $task.GetAwaiter().GetResult()

            Write-Host "VSCode attached to pipe. Sending back confirmation byte to debug session."
            $task = $stream.WriteAsync([byte[]]@(0), 0, 1)
            while (-not $task.AsyncWaitHandle.WaitOne(300)) {}
            $null = $task.GetAwaiter().GetResult()

            # FIXME: SSH socket forwarding won't break the copy so need to find out how to detect that.
            Write-Host "Starting socket <-> pipe streaming"
            $writeTask = $pipe.CopyToAsync($stream)
            $readTask = $stream.CopyToAsync($pipe)

            Write-Host "Waiting for startDebugging attach response to arrive"
            while (-not $attachTask.AsyncWaitHandle.WaitOne(300)) { }
            $null = $attachTask.GetAwaiter().GetResult()

            $task = [Task]::WhenAny($writeTask, $readTask)
            while (-not $task.AsyncWaitHandle.WaitOne(300)) {}
            $finishedTask = $task.GetAwaiter().GetResult()

            if ($finishedTask -eq $writeTask) {
                Write-Host "VSCode disconnected from Pipe $pipeName and RunspaceId $runspaceId"
            }
            elseif ($finishedTask -eq $readTask) {
                Write-Host "Socket disconnected from Pipe $pipeName and RunspaceId $runspaceId"
            }
            else {
                Write-Host "Unknown task finished for Pipe $pipeName and RunspaceId $runspaceId"
            }
        }
        finally {
            ${pipe}?.Dispose()
            ${stream}?.Dispose()
            $client.Dispose()
        }
    }
}
finally {
    if (Test-Path -LiteralPath $configFile) {
        Remove-Item -LiteralPath $configFile -Force
    }
    ${sock}?.Dispose()
}

#!/usr/bin/env python
# PYTHON_ARGCOMPLETE_OK

from __future__ import annotations

import argparse
import contextlib
import json
import os
import pathlib
import re
import subprocess
import sys
import typing as t

HAS_ARGCOMPLETE = False
try:
    import argcomplete

    HAS_ARGCOMPLETE = True
except ImportError:
    pass

_PORT_ALLOC_PATTERN = re.compile(r'Allocated port (\d+) for remote forward to .*')

def _parse_args(args: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Ansible Module Debugging CLI")

    parser.add_argument(
        'hostname',
        nargs=1,
        type=str,
        help="The SSH hostname to forward the debug ports to. This should include username and/or port if necessary.",
    )

    if HAS_ARGCOMPLETE:
        argcomplete.autocomplete(parser)

    return parser.parse_args(args)


@contextlib.contextmanager
def forward_ports(
    hostname: str,
    local_port: int,
) -> t.Generator[int]:
    ssh_cmd = [
        'ssh',
        '-v',  # Verbose output to capture port allocation
        '-NT',  # Don't execute remote command, disable PTY allocation
        '-o', 'ExitOnForwardFailure=yes',
        '-R', f'localhost:0:localhost:{local_port}',
        hostname
    ]

    # Start SSH process with stderr captured for debug output
    process = subprocess.Popen(
        ssh_cmd,
        stderr=subprocess.STDOUT,
        stdout=subprocess.PIPE,
        text=True,
    )
    try:
        assert process.stdout is not None

        proc_out = []
        port = 0
        for line in process.stdout:
            if match := _PORT_ALLOC_PATTERN.match(line):
                port = int(match.group(1))
                break

            proc_out.append(line)

        if port == 0:
            raise Exception(f"SSH failed to forward ports. Output:\n{'\n'.join(proc_out)}")

        yield port

    finally:
        process.terminate()


def main() -> None:
    args = _parse_args(sys.argv[1:])

    config_path = pathlib.Path.home() / '.ansible' / 'test' / 'debugging' / 'pwsh_debugger.json'
    if not config_path.exists():
        raise Exception(f"PowerShell debugger configuration file not found. Please ensure it exists at {config_path.absolute()}")

    debug_info = json.loads(config_path.read_text())
    local_port = debug_info.get('port', 0)
    if not local_port:
        raise Exception("No local port specified in the PowerShell debugger configuration file.")

    with forward_ports(args.hostname[0], local_port=local_port) as remote_port:
        debug_options = {
            'wait': False,
            'host': 'localhost',
            'port': remote_port,
        }
        new_env = os.environ.copy() | {
            '_ANSIBLE_ANSIBALLZ_PWSH_DEBUGGER_CONFIG': json.dumps(debug_options),
        }
        res = subprocess.run('/bin/bash', env=new_env, check=False)
        sys.exit(res.returncode)


if __name__ == '__main__':
    main()