From 1720601b5541bc8331375912ec0f8273d3e847ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=8D=E5=BF=98=E5=88=9D=E5=BF=83?= <1624717079@qq.com> Date: Sat, 14 Sep 2024 12:01:56 +0800 Subject: [PATCH 1/6] [Fix] Fix ssh tunnel always return 127.0.0.1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 修复ssh隧道始终返回127.0.0.1，改为动态获取本地地址 --- sql/utils/ssh_tunnel.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/sql/utils/ssh_tunnel.py b/sql/utils/ssh_tunnel.py index 2946eb5f1f..32d4f6dbec 100644 --- a/sql/utils/ssh_tunnel.py +++ b/sql/utils/ssh_tunnel.py @@ -8,6 +8,7 @@ from sshtunnel import SSHTunnelForwarder from paramiko import RSAKey import io +import socket class SSHConnection(object): @@ -55,13 +56,20 @@ def __init__( ) self.server.start() + # 动态获取本地IP + self.local_ip = self.get_local_ip() + def __del__(self): self.server.close() + def get_local_ip(self): + """ + 动态获取本地IP地址 + """ + return socket.gethostbyname(socket.gethostname()) + def get_ssh(self): """ - 获取ssh映射的端口 - :param request: - :return: + 获取ssh映射的本地IP和端口 """ - return "127.0.0.1", self.server.local_bind_port + return self.local_ip, self.server.local_bind_port From 9b4c44bba9687d88d0b77fa2fc92f97af7f82107 Mon Sep 17 00:00:00 2001 From: Leo Q Date: Sat, 14 Sep 2024 12:14:57 +0800 Subject: [PATCH 2/6] add tests for ssh tunnel --- sql/utils/tests.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/sql/utils/tests.py b/sql/utils/tests.py index 6ece20636b..f88d6d8a3d 100644 --- a/sql/utils/tests.py +++ b/sql/utils/tests.py @@ -39,6 +39,7 @@ from sql.utils.execute_sql import execute, execute_callback from sql.utils.tasks import add_sql_schedule, del_schedule, task_info from sql.utils.data_masking import data_masking, brute_mask, simple_column_mask +from sql.utils.ssh_tunnel import SSHConnection User = Users __author__ = "hhyo" @@ -1512,3 +1513,8 @@ def test_auth_group_users(self): auth_group_names=[self.agp.name], group_id=self.rgp1.group_id ) self.assertIn(self.user, users) + +# ssh tunnel tests +def test_get_local_ip(): + tunnel = SSHConnection("mysql", 3306, "tunnel", 22, "root", "password", "", "") + assert tunnel.get_local_ip() != "127.0.0.1" From 9a9b3227423e7caacee1ce3e0278ff3381009ef1 Mon Sep 17 00:00:00 2001 From: Leo Q Date: Sat, 14 Sep 2024 12:16:27 +0800 Subject: [PATCH 3/6] Update tests.py --- sql/utils/tests.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sql/utils/tests.py b/sql/utils/tests.py index f88d6d8a3d..3f3e9fa3a6 100644 --- a/sql/utils/tests.py +++ b/sql/utils/tests.py @@ -1516,5 +1516,5 @@ def test_auth_group_users(self): # ssh tunnel tests def test_get_local_ip(): - tunnel = SSHConnection("mysql", 3306, "tunnel", 22, "root", "password", "", "") - assert tunnel.get_local_ip() != "127.0.0.1" + tunnel = SSHConnection("mysql", 3306, "tunnel", 22, "root", "password", "", "") + assert tunnel.get_local_ip() != "127.0.0.1" From 20ba7eccf41f41a996817788db61adb03c619b62 Mon Sep 17 00:00:00 2001 From: Leo Q Date: Sat, 14 Sep 2024 13:37:02 +0800 Subject: [PATCH 4/6] add tests, mock, assertion --- sql/utils/tests.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/sql/utils/tests.py b/sql/utils/tests.py index 3f3e9fa3a6..d61c0d1582 100644 --- a/sql/utils/tests.py +++ b/sql/utils/tests.py @@ -9,6 +9,7 @@ import datetime import json from unittest.mock import patch, MagicMock +from pytest_mock import MockerFixture from django.conf import settings from django.contrib.auth.models import Permission, Group @@ -1514,7 +1515,12 @@ def test_auth_group_users(self): ) self.assertIn(self.user, users) + # ssh tunnel tests -def test_get_local_ip(): +def test_get_local_ip(mocker: MockerFixture): + from sql.utils.ssh_tunnel import SSHTunnelForwarder + + mock_start = mocker.patch.object(SSHTunnelForwarder, "start") tunnel = SSHConnection("mysql", 3306, "tunnel", 22, "root", "password", "", "") assert tunnel.get_local_ip() != "127.0.0.1" + assert mock_start.assert_called_once() From 12e1bddb03d17c1d837a48f3071e9ca37daab630 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=8D=E5=BF=98=E5=88=9D=E5=BF=83?= <1624717079@qq.com> Date: Sat, 14 Sep 2024 15:21:38 +0800 Subject: [PATCH 5/6] [Fix] Fix ssh tunnel always return 127.0.0.1 --- sql/utils/ssh_tunnel.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/sql/utils/ssh_tunnel.py b/sql/utils/ssh_tunnel.py index 32d4f6dbec..d0b9bce626 100644 --- a/sql/utils/ssh_tunnel.py +++ b/sql/utils/ssh_tunnel.py @@ -63,10 +63,14 @@ def __del__(self): self.server.close() def get_local_ip(self): - """ - 动态获取本地IP地址 - """ - return socket.gethostbyname(socket.gethostname()) + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + try: + # 连接一个公网地址（不需要实际发送数据） + s.connect(('8.8.8.8', 80)) + ip = s.getsockname()[0] + finally: + s.close() + return ip def get_ssh(self): """ From 2659223d3a7286e1e048d60bc1a523c4f3ac2347 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=8D=E5=BF=98=E5=88=9D=E5=BF=83?= <1624717079@qq.com> Date: Sat, 14 Sep 2024 15:53:57 +0800 Subject: [PATCH 6/6] add tests for ssh tunnel --- sql/utils/tests.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/sql/utils/tests.py b/sql/utils/tests.py index d61c0d1582..b626470e2f 100644 --- a/sql/utils/tests.py +++ b/sql/utils/tests.py @@ -9,7 +9,6 @@ import datetime import json from unittest.mock import patch, MagicMock -from pytest_mock import MockerFixture from django.conf import settings from django.contrib.auth.models import Permission, Group @@ -1517,10 +1516,6 @@ def test_auth_group_users(self): # ssh tunnel tests -def test_get_local_ip(mocker: MockerFixture): - from sql.utils.ssh_tunnel import SSHTunnelForwarder - - mock_start = mocker.patch.object(SSHTunnelForwarder, "start") - tunnel = SSHConnection("mysql", 3306, "tunnel", 22, "root", "password", "", "") - assert tunnel.get_local_ip() != "127.0.0.1" - assert mock_start.assert_called_once() +def test_get_local_ip(): + tunnel = SSHConnection("mysql", 3306, "tunnel", 22, "root", "password", "", "") + assert tunnel.get_local_ip() != "127.0.0.1"