From b518f228ded7b12d3443374d24922175d2c1bcd5 Mon Sep 17 00:00:00 2001
From: Lonny Wong <lonnywong@qq.com>
Date: Sun, 26 May 2024 09:21:51 +0800
Subject: [PATCH] support using putty pageant as ssh agent

Based on #111 by abakum <koka.abakum@gmail.com>
---
 go.mod                |  1 +
 go.sum                |  2 ++
 tssh/agent.go         |  5 +----
 tssh/agent_unix.go    |  4 +++-
 tssh/agent_windows.go | 18 +++++++++++++++++-
 5 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/go.mod b/go.mod
index da5744f..b3fe692 100644
--- a/go.mod
+++ b/go.mod
@@ -20,6 +20,7 @@ require (
 	github.com/stretchr/testify v1.8.4
 	github.com/trzsz/go-arg v1.5.3
 	github.com/trzsz/iterm2 v0.1.2
+	github.com/trzsz/pageant v0.1.0
 	github.com/trzsz/promptui v0.10.7
 	github.com/trzsz/ssh_config v1.3.6
 	github.com/trzsz/trzsz-go v1.1.8-0.20240525015006-6424386a6738
diff --git a/go.sum b/go.sum
index c934de6..65760a2 100644
--- a/go.sum
+++ b/go.sum
@@ -97,6 +97,8 @@ github.com/trzsz/go-arg v1.5.3 h1:eIDwDEmvSahtr5HpQOLrSa+YMqWQQ0H20xx60XgXQJw=
 github.com/trzsz/go-arg v1.5.3/go.mod h1:IC6Z/FiVH7uYvcbp1/gJhDYCFPS/GkL0APYakVvgY4I=
 github.com/trzsz/iterm2 v0.1.2 h1:VwfLzr2fKeaLf+p4tS0ms+kqdiQQxVLbTJUoyuQXmK8=
 github.com/trzsz/iterm2 v0.1.2/go.mod h1:PMI+3JcT7J9D0T6e3mOWv8ICYdrrNZwuge3Tm7zDLws=
+github.com/trzsz/pageant v0.1.0 h1:Ux5VSA6IMwahrMQ+8BeCnDKh3lO/OPkCQgBEpsQ379s=
+github.com/trzsz/pageant v0.1.0/go.mod h1:OMX42Kisk5o8eVDR/UCHmNPaeRf8MnJ33muUTqSwDO4=
 github.com/trzsz/promptui v0.10.7 h1:77uBrmsIPYYJS/9n+zwFRhwOz82EKXkkdjOiWSEUPpk=
 github.com/trzsz/promptui v0.10.7/go.mod h1:9dp59ixe32qBV9GjDxQ1PDWwbzHjTzveZenQwEoVHbg=
 github.com/trzsz/ssh_config v1.3.6 h1:+LBCg2uzhAgw2s19yqeUdD4YwW2z4kvlsXtKB6zDjmQ=
diff --git a/tssh/agent.go b/tssh/agent.go
index faa8894..7906d8c 100644
--- a/tssh/agent.go
+++ b/tssh/agent.go
@@ -53,10 +53,7 @@ func getAgentAddr(args *sshArgs, param *sshParam) (string, error) {
 	if addr := os.Getenv("SSH_AUTH_SOCK"); addr != "" {
 		return resolveHomeDir(addr), nil
 	}
-	if addr := defaultAgentAddr; addr != "" && isFileExist(addr) {
-		return addr, nil
-	}
-	return "", nil
+	return getDefaultAgentAddr()
 }
 
 func getAgentClient(args *sshArgs, param *sshParam) agent.ExtendedAgent {
diff --git a/tssh/agent_unix.go b/tssh/agent_unix.go
index 0bb5bc6..dacb927 100644
--- a/tssh/agent_unix.go
+++ b/tssh/agent_unix.go
@@ -31,7 +31,9 @@ import (
 	"time"
 )
 
-const defaultAgentAddr = ""
+func getDefaultAgentAddr() (string, error) {
+	return "", nil
+}
 
 func dialAgent(addr string) (net.Conn, error) {
 	return net.DialTimeout("unix", addr, time.Second)
diff --git a/tssh/agent_windows.go b/tssh/agent_windows.go
index 7e35581..64d4009 100644
--- a/tssh/agent_windows.go
+++ b/tssh/agent_windows.go
@@ -29,11 +29,27 @@ import (
 	"time"
 
 	"github.com/Microsoft/go-winio"
+	"github.com/trzsz/pageant"
 )
 
-const defaultAgentAddr = `\\.\pipe\openssh-ssh-agent`
+const kSshAgentAddr = `\\.\pipe\openssh-ssh-agent`
+
+const kPageantFakeAddr = "using_pageant_as_ssh_agent"
+
+func getDefaultAgentAddr() (string, error) {
+	if isFileExist(kSshAgentAddr) {
+		return kSshAgentAddr, nil
+	}
+	if pageant.PageantAvailable() {
+		return kPageantFakeAddr, nil
+	}
+	return "", nil
+}
 
 func dialAgent(addr string) (net.Conn, error) {
+	if addr == kPageantFakeAddr {
+		return pageant.NewPageantConn()
+	}
 	timeout := time.Second
 	return winio.DialPipe(addr, &timeout)
 }