feat: support SSH IdentityAgent config directive (#1630)

Author: wowi42

src/pyinfra/connectors/ssh.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from random import uniform
 from shutil import which
 from socket import error as socket_error, gaierror
@@ -308,10 +309,23 @@ def _retry_paramiko_agent_keys(
         if not kwargs.get("allow_agent"):
             return False
 
+        # Honor IdentityAgent from SSH config
+        identity_agent = getattr(self.client, "identity_agent", None)
+        old_auth_sock = os.environ.get("SSH_AUTH_SOCK")
+        if isinstance(identity_agent, str):
+            os.environ["SSH_AUTH_SOCK"] = identity_agent
+        else:
+            identity_agent = None
         try:
             agent_keys = list(Agent().get_keys())
         except Exception:
             return False
+        finally:
+            if identity_agent:
+                if old_auth_sock is not None:
+                    os.environ["SSH_AUTH_SOCK"] = old_auth_sock
+                else:
+                    os.environ.pop("SSH_AUTH_SOCK", None)
 
         if not agent_keys:
             return False

src/pyinfra/connectors/sshuserclient/client.py

@@ -3,6 +3,7 @@
 source has now vanished (https://github.com/tobald/sshuserclient).
 """
 
+import os
 from os import path
 
 from gevent.lock import BoundedSemaphore
@@ -166,6 +167,7 @@ def connect(  # type: ignore[override]
             missing_host_key_policy,
             host_keys_files,
             keep_alive,
+            identity_agent,
         ) = self.parse_config(
             hostname,
             kwargs,
@@ -188,7 +190,21 @@ def connect(  # type: ignore[override]
             config.update(_pyinfra_ssh_paramiko_connect_kwargs)
 
         self._ssh_config = config
-        super().connect(hostname, **config)
+        self.identity_agent = identity_agent
+
+        # Honor IdentityAgent from SSH config by temporarily setting SSH_AUTH_SOCK
+        # so Paramiko's Agent class connects to the correct socket.
+        old_auth_sock = os.environ.get("SSH_AUTH_SOCK")
+        if identity_agent:
+            os.environ["SSH_AUTH_SOCK"] = identity_agent
+        try:
+            super().connect(hostname, **config)
+        finally:
+            if identity_agent:
+                if old_auth_sock is not None:
+                    os.environ["SSH_AUTH_SOCK"] = old_auth_sock
+                else:
+                    os.environ.pop("SSH_AUTH_SOCK", None)
 
         if _pyinfra_ssh_forward_agent is not None:
             forward_agent = _pyinfra_ssh_forward_agent
@@ -225,6 +241,7 @@ def parse_config(
 
         keep_alive = 0
         forward_agent = False
+        identity_agent = None
         missing_host_key_policy = get_missing_host_key_policy(strict_host_key_checking)
         host_keys_files = (path.expanduser("~/.ssh/known_hosts"),)
 
@@ -237,6 +254,7 @@ def parse_config(
                 missing_host_key_policy,
                 host_keys_files,
                 keep_alive,
+                identity_agent,
             )
 
         host_config = ssh_config.lookup(hostname)
@@ -269,6 +287,11 @@ def parse_config(
         if "serveraliveinterval" in host_config:
             keep_alive = int(host_config["serveraliveinterval"])
 
+        if "identityagent" in host_config:
+            agent_path = host_config["identityagent"]
+            if agent_path.lower() != "none":
+                identity_agent = path.expanduser(agent_path)
+
         if "proxycommand" in host_config:
             cfg["sock"] = ProxyCommand(host_config["proxycommand"])
 
@@ -294,7 +317,15 @@ def parse_config(
                 sock = c.gateway(hostname, cfg["port"], target, target_config["port"])
             cfg["sock"] = sock
 
-        return hostname, cfg, forward_agent, missing_host_key_policy, host_keys_files, keep_alive
+        return (
+            hostname,
+            cfg,
+            forward_agent,
+            missing_host_key_policy,
+            host_keys_files,
+            keep_alive,
+            identity_agent,
+        )
 
     @staticmethod
     def derive_shorthand(ssh_config, host_string):

tests/test_connectors/test_sshuserclient.py

@@ -46,6 +46,18 @@
     UserKnownHostsFile ~/.ssh/known_hosts ~/.ssh/known_hosts.infra ~/.ssh/known_hosts.webservers
 """
 
+SSH_CONFIG_IDENTITY_AGENT = """
+Host 10.0.0.1
+    User agentuser
+    IdentityAgent ~/Library/Group Containers/2BUA8C4S2C.com.1password/t/agent.sock
+"""
+
+SSH_CONFIG_IDENTITY_AGENT_NONE = """
+Host 10.0.0.2
+    User agentuser
+    IdentityAgent none
+"""
+
 BAD_SSH_CONFIG_DATA = """
 &
 """
@@ -90,13 +102,20 @@ def setUp(self):
     def test_load_ssh_config_no_exist(self):
         client = SSHClient()
 
-        _, config, forward_agent, missing_host_key_policy, host_keys_file, keep_alive = (
-            client.parse_config(
-                "127.0.0.1",
-            )
+        (
+            _,
+            config,
+            forward_agent,
+            missing_host_key_policy,
+            host_keys_file,
+            keep_alive,
+            identity_agent,
+        ) = client.parse_config(
+            "127.0.0.1",
         )
 
         assert config.get("port") == 22
+        assert identity_agent is None
 
 
 @patch(
@@ -140,10 +159,16 @@ def setUp(self):
     def test_load_ssh_config(self):
         client = SSHClient()
 
-        _, config, forward_agent, missing_host_key_policy, host_keys_file, keep_alive = (
-            client.parse_config(
-                "127.0.0.1",
-            )
+        (
+            _,
+            config,
+            forward_agent,
+            missing_host_key_policy,
+            host_keys_file,
+            keep_alive,
+            identity_agent,
+        ) = client.parse_config(
+            "127.0.0.1",
         )
 
         assert config.get("key_filename") == ["/id_rsa", "/id_rsa2"]
@@ -153,6 +178,7 @@ def test_load_ssh_config(self):
         assert forward_agent is False
         assert isinstance(missing_host_key_policy, AskPolicy)
         assert host_keys_file == ("~/.ssh/known_hosts",)  # OpenSSH default
+        assert identity_agent is None
 
         (
             _,
@@ -161,6 +187,7 @@ def test_load_ssh_config(self):
             missing_host_key_policy,
             host_keys_file,
             keep_alive,
+            identity_agent,
         ) = client.parse_config("192.168.1.1")
 
         assert other_config.get("username") == "otheruser"
@@ -177,9 +204,15 @@ def test_load_ssh_config_inline_comments(self):
         """Test that inline comments are stripped from SSH config values (issue #1568)."""
         client = SSHClient()
 
-        _, config, forward_agent, missing_host_key_policy, host_keys_file, keep_alive = (
-            client.parse_config("127.0.0.1")
-        )
+        (
+            _,
+            config,
+            forward_agent,
+            missing_host_key_policy,
+            host_keys_file,
+            keep_alive,
+            identity_agent,
+        ) = client.parse_config("127.0.0.1")
 
         assert config.get("key_filename") == ["/id_rsa"]
         assert config.get("username") == "testuser"
@@ -206,6 +239,7 @@ def test_load_ssh_config_multiple_known_hosts(self):
             missing_host_key_policy,
             host_keys_files,
             keep_alive,
+            identity_agent,
         ) = client.parse_config("192.168.1.3")
 
         # Verify multiple known hosts files are parsed as a tuple
@@ -215,6 +249,34 @@ def test_load_ssh_config_multiple_known_hosts(self):
             "~/.ssh/known_hosts.webservers",
         )
 
+    @patch(
+        "pyinfra.connectors.sshuserclient.client.open",
+        mock_open(read_data=SSH_CONFIG_IDENTITY_AGENT),
+        create=True,
+    )
+    def test_load_ssh_config_identity_agent(self):
+        """Test that IdentityAgent is parsed from SSH config."""
+        client = SSHClient()
+
+        _, config, _, _, _, _, identity_agent = client.parse_config("10.0.0.1")
+
+        assert config.get("username") == "agentuser"
+        assert identity_agent == "~/Library/Group Containers/2BUA8C4S2C.com.1password/t/agent.sock"
+
+    @patch(
+        "pyinfra.connectors.sshuserclient.client.open",
+        mock_open(read_data=SSH_CONFIG_IDENTITY_AGENT_NONE),
+        create=True,
+    )
+    def test_load_ssh_config_identity_agent_none(self):
+        """Test that IdentityAgent set to 'none' is ignored."""
+        client = SSHClient()
+
+        _, config, _, _, _, _, identity_agent = client.parse_config("10.0.0.2")
+
+        assert config.get("username") == "agentuser"
+        assert identity_agent is None
+
     @patch(
         "pyinfra.connectors.sshuserclient.client.open",
         mock_open(read_data=BAD_SSH_CONFIG_DATA),
@@ -262,7 +324,7 @@ def test_load_ssh_config_proxyjump(self, fake_gateway, fake_ssh_connect):
         client = SSHClient()
 
         # Load the SSH config with ProxyJump configured
-        _, config, forward_agent, _, _, _ = client.parse_config(
+        _, config, forward_agent, _, _, _, _ = client.parse_config(
             "192.168.1.2",
             {"port": 1022},
             ssh_config_file="other_file",