Skip to content

Commit 0f7c999

Browse files
committed
Timeout DB2 detector if it takes too long (Yelp#214)
* Timeout DB2 detector if it takes too long * Trying to fix tests * Use multiprocessing instead of signal * Faster tests * whitespace and comments * Print 'DB2Detector timed out.' * Use SQL_ATTR_LOGIN_TIMEOUT * Fix tests * ConnectTimeout * Error message * Addressed @xianjun comments * Timeout is UNVERIFIED, other exceptions are VERIFIED_FALSE * Addressed @xianjun comment
1 parent 70f0ce0 commit 0f7c999

File tree

2 files changed

+34
-10
lines changed

2 files changed

+34
-10
lines changed

detect_secrets/plugins/db2.py

+13-6
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ class DB2Detector(RegexBasedDetector):
5757
r'*(?:.\[A-Za-z0-9]|[A-Za-z0-9][A-Za-z0-9\-]*[A-Za-z0-9]))'
5858
)
5959

60-
def verify(self, token, content, potential_secret=None):
60+
def verify(self, token, content, potential_secret=None, timeout=5):
6161

6262
username_matches = get_other_factor(
6363
content, self.username_keyword_regex,
@@ -93,7 +93,7 @@ def verify(self, token, content, potential_secret=None):
9393
for port in port_matches: # pragma: no cover
9494
for hostname in hostname_matches: # pragma: no cover
9595
verify_result = verify_db2_credentials(
96-
database, hostname, port, username, token,
96+
database, hostname, port, username, token, timeout,
9797
)
9898
if verify_result == VerifiedResult.VERIFIED_TRUE:
9999
potential_secret.other_factors['database'] = database
@@ -105,24 +105,31 @@ def verify(self, token, content, potential_secret=None):
105105
return VerifiedResult.VERIFIED_FALSE
106106

107107

108-
def verify_db2_credentials(database, hostname, port, username, password): # pragma: no cover
108+
def verify_db2_credentials(
109+
database, hostname, port, username, password, timeout=5,
110+
): # pragma: no cover
109111
try:
110112
conn_str = 'database={database};hostname={hostname};port={port};' + \
111-
'protocol=tcpip;uid={username};pwd={password}'
113+
'protocol=tcpip;uid={username};pwd={password};' + \
114+
'ConnectTimeout={timeout}'
112115
conn_str = conn_str.format(
113116
database=database,
114117
hostname=hostname,
115118
port=port,
116119
username=username,
117120
password=password,
121+
timeout=timeout,
118122
)
119123
ibm_db_conn = ibm_db.connect(conn_str, '', '')
120124
if ibm_db_conn:
121125
return VerifiedResult.VERIFIED_TRUE
122126
else:
123127
return VerifiedResult.VERIFIED_FALSE
124-
except Exception:
125-
return VerifiedResult.UNVERIFIED
128+
except Exception as e:
129+
if 'Timeout' in str(e):
130+
return VerifiedResult.UNVERIFIED
131+
else:
132+
return VerifiedResult.VERIFIED_FALSE
126133

127134

128135
def get_other_factor(content, factor_keyword_regex, factor_regex):

tests/plugins/db2_test.py

+21-4
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
DB2_HOSTNAME = 'fake.host.name'
2020
DB2_DATABASE = 'fake_database'
2121
DB2_CONN_STRING = 'database={DB2_DATABASE};hostname={DB2_HOSTNAME};port={DB2_PORT};' + \
22-
'protocol=tcpip;uid={DB2_USER};pwd={DB2_PASSWORD}'
22+
'protocol=tcpip;uid={DB2_USER};pwd={DB2_PASSWORD};ConnectTimeout=5'
2323
DB2_CONN_STRING = DB2_CONN_STRING.format(
2424
DB2_DATABASE=DB2_DATABASE,
2525
DB2_HOSTNAME=DB2_HOSTNAME,
@@ -59,7 +59,7 @@ def test_analyze_string(self, token, payload, should_flag):
5959
assert list(output.keys())[0].secret == token
6060

6161
@patch('detect_secrets.plugins.db2.ibm_db.connect')
62-
def test_verify_invalid_secret(self, mock_db2_connect):
62+
def test_verify_invalid_connect_returns_none(self, mock_db2_connect):
6363
mock_db2_connect.return_value = None
6464

6565
potential_secret = PotentialSecret('test db2', 'test filename', DB2_PASSWORD)
@@ -75,6 +75,23 @@ def test_verify_invalid_secret(self, mock_db2_connect):
7575

7676
mock_db2_connect.assert_called_with(DB2_CONN_STRING, '', '')
7777

78+
@patch('detect_secrets.plugins.db2.ibm_db.connect')
79+
def test_verify_invalid_connect_throws_exception(self, mock_db2_connect):
80+
mock_db2_connect.side_effect = Exception('oops')
81+
82+
potential_secret = PotentialSecret('test db2', 'test filename', DB2_PASSWORD)
83+
assert DB2Detector().verify(
84+
DB2_PASSWORD,
85+
'''user={},
86+
password={},
87+
database={},
88+
host={},
89+
port={}'''.format(DB2_USER, DB2_PASSWORD, DB2_DATABASE, DB2_HOSTNAME, DB2_PORT),
90+
potential_secret,
91+
) == VerifiedResult.VERIFIED_FALSE
92+
93+
mock_db2_connect.assert_called_with(DB2_CONN_STRING, '', '')
94+
7895
@patch('detect_secrets.plugins.db2.ibm_db.connect')
7996
def test_verify_valid_secret(self, mock_db2_connect):
8097
mock_db2_connect.return_value = MagicMock()
@@ -145,8 +162,8 @@ def test_verify_from_url(self, mock_db2_connect):
145162
mock_db2_connect.assert_called_with(DB2_CONN_STRING, '', '')
146163

147164
@patch('detect_secrets.plugins.db2.ibm_db.connect')
148-
def test_verify_unverified_secret(self, mock_db2_connect):
149-
mock_db2_connect.side_effect = Exception('oops')
165+
def test_verify_times_out(self, mock_db2_connect):
166+
mock_db2_connect.side_effect = Exception('Timeout')
150167

151168
potential_secret = PotentialSecret('test db2', 'test filename', DB2_PASSWORD)
152169
assert DB2Detector().verify(

0 commit comments

Comments
 (0)