Skip to content

Commit b54be42

Browse files
committed
Refactor SSL socket handling to use wrap_socket_to_ssl method
Update `test_ssl.py`, `socket.py`, and `ssl.py` to replace direct calls to `ssl.wrap_socket` with a new `wrap_socket_to_ssl` function. Adjust test cases accordingly to ensure proper integration and functionality. This improves code organization and makes SSL socket wrapping more explicit.Refactor SSL socket wrapping logic Moved SSL socket wrapping to a dedicated function `wrap_socket_to_ssl` within the `ssl.py` module, improving code modularity and readability. Updated the imports and references across the codebase to reflect this change, enhancing maintainability and reducing redundancy.
1 parent f65f9bc commit b54be42

File tree

4 files changed

+17
-33
lines changed

4 files changed

+17
-33
lines changed

nettacker/core/lib/socket.py

+4-10
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,11 @@
66
import re
77
import select
88
import socket
9-
import ssl
109
import struct
1110
import time
1211

1312
from nettacker.core.lib.base import BaseEngine, BaseLibrary
13+
from nettacker.core.lib.ssl import wrap_socket_to_ssl
1414
from nettacker.core.utils.common import reverse_and_regex_condition
1515

1616
log = logging.getLogger(__name__)
@@ -21,21 +21,15 @@ def create_tcp_socket(host, port, timeout):
2121
socket_connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
2222
socket_connection.settimeout(timeout)
2323
socket_connection.connect((host, port))
24-
ssl_flag = False
2524
except ConnectionRefusedError:
2625
return None
2726

2827
try:
29-
socket_connection = ssl.wrap_socket(socket_connection)
30-
ssl_flag = True
28+
return wrap_socket_to_ssl(socket_connection), True
3129
except Exception:
32-
socket_connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
33-
socket_connection.settimeout(timeout)
34-
socket_connection.connect((host, port))
35-
# finally:
36-
# socket_connection.shutdown()
30+
pass
3731

38-
return socket_connection, ssl_flag
32+
return socket_connection, False
3933

4034

4135
class SocketLibrary(BaseLibrary):

nettacker/core/lib/ssl.py

+8-18
Original file line numberDiff line numberDiff line change
@@ -107,24 +107,10 @@ def test_single_cipher(host, port, cipher, timeout):
107107
return supported_ciphers, False
108108

109109

110-
def create_tcp_socket(host, port, timeout):
111-
try:
112-
socket_connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
113-
socket_connection.settimeout(timeout)
114-
socket_connection.connect((host, port))
115-
ssl_flag = False
116-
except ConnectionRefusedError:
117-
return None
118-
119-
try:
120-
socket_connection = ssl.wrap_socket(socket_connection)
121-
ssl_flag = True
122-
except Exception:
123-
socket_connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
124-
socket_connection.settimeout(timeout)
125-
socket_connection.connect((host, port))
126-
127-
return socket_connection, ssl_flag
110+
# nosec B504 # ignore SSL/TLS version warning
111+
def wrap_socket_to_ssl(socket_connection):
112+
context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
113+
return context.wrap_socket(socket_connection)
128114

129115

130116
def get_cert_info(cert):
@@ -154,6 +140,8 @@ def get_cert_info(cert):
154140

155141
class SslLibrary(BaseLibrary):
156142
def ssl_certificate_scan(self, host, port, timeout):
143+
from nettacker.core.lib.socket import create_tcp_socket
144+
157145
tcp_socket = create_tcp_socket(host, port, timeout)
158146
if tcp_socket is None:
159147
return None
@@ -175,6 +163,8 @@ def ssl_certificate_scan(self, host, port, timeout):
175163
return scan_info
176164

177165
def ssl_version_and_cipher_scan(self, host, port, timeout):
166+
from nettacker.core.lib.socket import create_tcp_socket
167+
178168
tcp_socket = create_tcp_socket(host, port, timeout)
179169
if tcp_socket is None:
180170
return None

tests/core/lib/test_socket.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ class Substeps:
114114

115115
class TestSocketMethod(TestCase):
116116
@patch("socket.socket")
117-
@patch("ssl.wrap_socket")
117+
@patch("nettacker.core.lib.socket.wrap_socket_to_ssl")
118118
def test_create_tcp_socket(self, mock_wrap, mock_socket):
119119
HOST = "example.com"
120120
PORT = 80

tests/core/lib/test_ssl.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
import ssl
22
from unittest.mock import patch
33

4+
from nettacker.core.lib.socket import create_tcp_socket
45
from nettacker.core.lib.ssl import (
56
SslEngine,
67
SslLibrary,
7-
create_tcp_socket,
88
is_weak_hash_algo,
99
is_weak_ssl_version,
1010
is_weak_cipher_suite,
@@ -153,7 +153,7 @@ class Substeps:
153153

154154
class TestSocketMethod(TestCase):
155155
@patch("socket.socket")
156-
@patch("ssl.wrap_socket")
156+
@patch("nettacker.core.lib.socket.wrap_socket_to_ssl")
157157
def test_create_tcp_socket(self, mock_wrap, mock_socket):
158158
HOST = "example.com"
159159
PORT = 80
@@ -167,7 +167,7 @@ def test_create_tcp_socket(self, mock_wrap, mock_socket):
167167

168168
@patch("nettacker.core.lib.ssl.is_weak_cipher_suite")
169169
@patch("nettacker.core.lib.ssl.is_weak_ssl_version")
170-
@patch("nettacker.core.lib.ssl.create_tcp_socket")
170+
@patch("nettacker.core.lib.socket.create_tcp_socket")
171171
def test_ssl_version_and_cipher_scan(self, mock_connection, mock_ssl_check, mock_cipher_check):
172172
library = SslLibrary()
173173
HOST = "example.com"
@@ -222,7 +222,7 @@ def test_ssl_version_and_cipher_scan(self, mock_connection, mock_ssl_check, mock
222222
},
223223
)
224224

225-
@patch("nettacker.core.lib.ssl.create_tcp_socket")
225+
@patch("nettacker.core.lib.socket.create_tcp_socket")
226226
@patch("nettacker.core.lib.ssl.is_weak_hash_algo")
227227
@patch("nettacker.core.lib.ssl.crypto.load_certificate")
228228
@patch("nettacker.core.lib.ssl.ssl.get_server_certificate")

0 commit comments

Comments
 (0)