Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support SSH proxy command #493

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Support SSH proxy command #493

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions dpdispatcher/contexts/ssh_context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
tar_compress=True,
look_for_keys=True,
execute_command=None,
proxy_command=None,
):
self.hostname = hostname
self.username = username
Expand All @@ -59,6 +60,7 @@
self.look_for_keys = look_for_keys
self.execute_command = execute_command
self._keyboard_interactive_auth = False
self.proxy_command = proxy_command
self._setup_ssh()

# @classmethod
Expand Down Expand Up @@ -141,12 +143,16 @@
# )
# assert(self.ssh.get_transport().is_active())
# transport = self.ssh.get_transport()
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.settimeout(self.timeout)
sock.connect((self.hostname, self.port))
if self.proxy_command is not None:
sock = paramiko.ProxyCommand(self.proxy_command)
sock.settimeout(self.timeout)

Check warning on line 148 in dpdispatcher/contexts/ssh_context.py

View check run for this annotation

Codecov / codecov/patch

dpdispatcher/contexts/ssh_context.py#L147-L148 

Added lines #L147 - L148 were not covered by tests
Comment on lines +146 to +148
Copy link
Contributor

@coderabbitai coderabbitai bot Sep 26, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Add unit tests for the proxy_command functionality

The newly added code handling the proxy_command is not covered by existing tests. To ensure reliability and prevent future regressions, please add unit tests that verify the behavior when a proxy_command is provided.

🧰 Tools
GitHub Check: codecov/patch

[warning] 147-148: dpdispatcher/contexts/ssh_context.py#L147-L148
Added lines #L147 - L148 were not covered by tests

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with the AI. A unit test should be added.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@njzjz, I'll proceed to add unit tests for the proxy_command functionality as requested. Here's a test case that verifies the behavior when a proxy_command is provided:

import unittest
from unittest.mock import patch, MagicMock
from dpdispatcher.contexts.ssh_context import SSHSession

class TestSSHSessionProxyCommand(unittest.TestCase):
    
    @patch('paramiko.ProxyCommand')
    @patch('socket.socket')
    def test_setup_ssh_with_proxy_command(self, mock_socket, mock_proxy_command):
        # Mock the ProxyCommand and socket behavior
        mock_proxy = MagicMock()
        mock_proxy_command.return_value = mock_proxy
        
        # Create an SSHSession with a proxy command
        session = SSHSession(
            hostname="example.com",
            username="user",
            proxy_command="proxy_command_example"
        )
        
        # Ensure that ProxyCommand is used instead of socket
        mock_proxy_command.assert_called_once_with("proxy_command_example")
        mock_proxy.settimeout.assert_called_once_with(session.timeout)
        mock_socket.assert_not_called()  # Ensure socket is not used

    @patch('paramiko.ProxyCommand')
    @patch('socket.socket')
    def test_setup_ssh_without_proxy_command(self, mock_socket, mock_proxy_command):
        # Create an SSHSession without a proxy command
        session = SSHSession(
            hostname="example.com",
            username="user"
        )
        
        # Ensure that socket is used instead of ProxyCommand
        mock_socket.assert_called_once_with(socket.AF_INET, socket.SOCK_STREAM)
        mock_socket().settimeout.assert_called_once_with(session.timeout)
        mock_socket().connect.assert_called_once_with(("example.com", 22))
        mock_proxy_command.assert_not_called()  # Ensure ProxyCommand is not used

if __name__ == '__main__':
    unittest.main()

Explanation:

  • Test 1 (test_setup_ssh_with_proxy_command): Verifies that when a proxy_command is provided, paramiko.ProxyCommand is used, and the standard socket connection is not invoked.
  • Test 2 (test_setup_ssh_without_proxy_command): Verifies that when no proxy_command is provided, the standard socket connection is used.

This should ensure that the proxy_command functionality is properly tested and prevent future regressions.

Let me know if you'd like any further adjustments!

else:
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.settimeout(self.timeout)
sock.connect((self.hostname, self.port))

# Make a Paramiko Transport object using the socket
ts = paramiko.Transport(sock)

Check failure on line 155 in dpdispatcher/contexts/ssh_context.py

View workflow job for this annotation

GitHub Actions / pyright 

Argument of type "ProxyCommand | socket" cannot be assigned to parameter "sock" of type "_SocketLike" in function "__init__"
  Type "ProxyCommand | socket" cannot be assigned to type "_SocketLike"
    Type "ProxyCommand" cannot be assigned to type "_SocketLike"
      "ProxyCommand" is incompatible with "str"
      "ProxyCommand" is incompatible with "_Addr"
      "ProxyCommand" is incompatible with "socket"
      "ProxyCommand" is incompatible with "Channel" (reportGeneralTypeIssues)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you ignore this line before python/typeshed#12697 is accepted?

ts.banner_timeout = 60
ts.auth_timeout = self.timeout + 20
ts.use_compression(compress=True)
Expand Down Expand Up @@ -343,6 +349,7 @@
"enable searching for discoverable private key files in ~/.ssh/"
)
doc_execute_command = "execute command after ssh connection is established."
doc_proxy_command = "The command of proxy."
ssh_remote_profile_args = [
Argument("hostname", str, optional=False, doc=doc_hostname),
Argument("username", str, optional=False, doc=doc_username),
Expand Down Expand Up @@ -391,6 +398,13 @@
default=None,
doc=doc_execute_command,
),
Argument(
"proxy_command",
str,
optional=True,
default=None,
doc=doc_proxy_command,
),
Comment on lines +401 to +407
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Ensure consistency in argument definitions

The Argument definition for proxy_command should align with the improved documentation string. Please update the doc parameter accordingly to reflect the clarified description.

]
ssh_remote_profile_format = Argument(
"ssh_session", dict, ssh_remote_profile_args
Expand Down
Loading