From 0785454a26099631ae035be55cf8e9a5a49be2d0 Mon Sep 17 00:00:00 2001
From: Jeff Forcier <jeff@bitprophet.org>
Date: Tue, 16 May 2023 23:09:44 -0400
Subject: Migrate rest of existing agent tests to new module

Also make them better.

Also fix AttributeError raising to not be py3.10 specific
---
 paramiko/agent.py |  3 ++-
 tests/agent.py    | 75 +++++++++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 72 insertions(+), 6 deletions(-)

diff --git a/paramiko/agent.py b/paramiko/agent.py
index 8486d6ca..9d215b85 100644
--- a/paramiko/agent.py
+++ b/paramiko/agent.py
@@ -55,6 +55,7 @@ ALGORITHM_FLAG_MAP = {
 }
 
 
+# TODO 4.0: rename all these - including making some of their methods public?
 class AgentSSH:
     def __init__(self):
         self._conn = None
@@ -470,7 +471,7 @@ class AgentKey(PKey):
         Proxy any un-implemented methods/properties to the inner_key.
         """
         if self.inner_key is None:  # nothing to proxy to
-            raise AttributeError(name=name, obj=self)
+            raise AttributeError(name)
         return getattr(self.inner_key, name)
 
     @property
diff --git a/tests/agent.py b/tests/agent.py
index 158a8b41..8e859289 100644
--- a/tests/agent.py
+++ b/tests/agent.py
@@ -1,4 +1,14 @@
-from paramiko import AgentKey
+from unittest.mock import Mock
+
+from pytest import mark, raises
+
+from paramiko import AgentKey, Message
+from paramiko.agent import (
+    SSH2_AGENT_SIGN_RESPONSE,
+    SSH_AGENT_RSA_SHA2_256,
+    SSH_AGENT_RSA_SHA2_512,
+    cSSH2_AGENTC_SIGN_REQUEST,
+)
 
 
 # AgentKey with no inner_key
@@ -17,15 +27,41 @@ class AgentKey_:
         key = AgentKey(None, b"secret!!!")
         assert str(key) == repr(key)
 
-    class fields:
-        """
-        _fields
-        """
+    class init:
+        def needs_at_least_two_arguments(self):
+            with raises(TypeError):
+                AgentKey()
+            with raises(TypeError):
+                AgentKey(None)
+
+        def sets_attributes_and_parses_blob(self):
+            agent = Mock()
+            blob = Message()
+            blob.add_string("bad-type")
+            key = AgentKey(agent=agent, blob=bytes(blob))
+            assert key.agent is agent
+            assert key.name == "bad-type"
+            assert key.blob == bytes(blob)
+            assert key.comment == ""  # default
+            # TODO: logger testing
+            assert key.inner_key is None  # no 'bad-type' algorithm
+
+        def comment_optional(self):
+            blob = Message()
+            blob.add_string("bad-type")
+            key = AgentKey(agent=Mock(), blob=bytes(blob), comment="hi!")
+            assert key.comment == "hi!"
 
+        def sets_inner_key_when_known_type(self, keys):
+            key = AgentKey(agent=Mock(), blob=bytes(keys.pkey))
+            assert key.inner_key == keys.pkey
+
+    class fields:
         def defaults_to_get_name_and_blob(self):
             key = _BareAgentKey(name="lol", blob=b"lmao")
             assert key._fields == ["lol", b"lmao"]
 
+        # TODO: pytest-relaxed is buggy (now?), this shows up under get_bits?
         def defers_to_inner_key_when_present(self, keys):
             key = AgentKey(agent=None, blob=keys.pkey.asbytes())
             assert key._fields == keys.pkey._fields
@@ -39,3 +75,32 @@ class AgentKey_:
         def defers_to_inner_key_when_present(self, keys):
             key = AgentKey(agent=None, blob=keys.pkey.asbytes())
             assert key.get_bits() == keys.pkey.get_bits()
+
+    @mark.parametrize(
+        "kwargs,expectation",
+        [
+            # No algorithm kwarg: no flags (bitfield -> 0 int)
+            (dict(), 0),
+            (dict(algorithm="rsa-sha2-256"), SSH_AGENT_RSA_SHA2_256),
+            (dict(algorithm="rsa-sha2-512"), SSH_AGENT_RSA_SHA2_512),
+        ],
+    )
+    def signing_data(self, kwargs, expectation):
+        class FakeAgent:
+            def _send_message(self, msg):
+                self._sent_message = msg
+                sig = Message()
+                sig.add_string("lol")
+                sig.rewind()
+                return SSH2_AGENT_SIGN_RESPONSE, sig
+
+        agent = FakeAgent()
+        key = AgentKey(agent, b"secret!!!")
+        result = key.sign_ssh_data(b"token", **kwargs)
+        assert result == b"lol"
+        msg = agent._sent_message
+        msg.rewind()
+        assert msg.get_byte() == cSSH2_AGENTC_SIGN_REQUEST
+        assert msg.get_string() == b"secret!!!"
+        assert msg.get_string() == b"token"
+        assert msg.get_int() == expectation
-- 
cgit v1.2.1