Repository navigation
Expand file tree
/
Copy pathtest_upgrade_sslcontext_tls.py
More file actions
137 lines (102 loc) · 3.94 KB
/
Copy pathtest_upgrade_sslcontext_tls.py
File metadata and controls
137 lines (102 loc) · 3.94 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import pytest
from core_codemods.upgrade_sslcontext_tls import UpgradeSSLContextTLS
from tests.codemods.base_codemod_test import BaseSemgrepCodemodTest
INSECURE_PROTOCOLS = [
"PROTOCOL_SSLv2",
"PROTOCOL_SSLv3",
"PROTOCOL_TLSv1",
"PROTOCOL_TLSv1_1",
"PROTOCOL_TLS", # This one isn't necessarily insecure, but it is deprecated
]
class TestUpgradeWeakTLS(BaseSemgrepCodemodTest):
codemod = UpgradeSSLContextTLS
@pytest.mark.parametrize("protocol", INSECURE_PROTOCOLS)
def test_upgrade_protocol_with_kwarg(self, tmpdir, protocol):
input_code = f"""import ssl
context = ssl.SSLContext(protocol=ssl.{protocol})
var = "hello"
"""
expected_output = """import ssl
context = ssl.SSLContext(protocol=ssl.PROTOCOL_TLS_CLIENT)
var = "hello"
"""
self.run_and_assert(tmpdir, input_code, expected_output)
@pytest.mark.parametrize("protocol", INSECURE_PROTOCOLS)
def test_upgrade_protocol_without_kwarg(self, tmpdir, protocol):
input_code = f"""import ssl
context = ssl.SSLContext(ssl.{protocol})
var = "hello"
"""
expected_output = """import ssl
context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
var = "hello"
"""
self.run_and_assert(tmpdir, input_code, expected_output)
@pytest.mark.parametrize("protocol", INSECURE_PROTOCOLS)
def test_upgrade_protocol_outside_sslcontext_dont_update(self, tmpdir, protocol):
input_code = f"""import ssl
foo(ssl.{protocol})
"""
expected_output = input_code
self.run_and_assert(tmpdir, input_code, expected_output)
@pytest.mark.parametrize("protocol", INSECURE_PROTOCOLS)
def test_upgrade_protocol_outside_call_dont_update(self, tmpdir, protocol):
input_code = f"""import ssl
if x == ssl.{protocol}:
print("whatever")
"""
expected_output = input_code
self.run_and_assert(tmpdir, input_code, expected_output)
@pytest.mark.skip("Need to add support for adding/rewriting imports")
@pytest.mark.parametrize("protocol", INSECURE_PROTOCOLS)
def test_upgrade_protocol_with_kwarg_direct_import(self, tmpdir, protocol):
input_code = f"""from ssl import {protocol}
context = ssl.SSLContext({protocol})
var = "hello"
"""
expected_output = """from ssl import PROTOCOL_TLS_CLIENT
context = ssl.SSLContext(PROTOCOL_TLS_CLIENT)
var = "hello"
"""
self.run_and_assert(tmpdir, input_code, expected_output)
@pytest.mark.parametrize("protocol", INSECURE_PROTOCOLS)
def test_upgrade_protocol_with_kwarg_import_alias(self, tmpdir, protocol):
input_code = f"""import ssl as whatever
context = whatever.SSLContext(protocol=whatever.{protocol})
var = "hello"
"""
expected_output = """import ssl as whatever
import ssl
context = whatever.SSLContext(protocol=ssl.PROTOCOL_TLS_CLIENT)
var = "hello"
"""
self.run_and_assert(tmpdir, input_code, expected_output)
def test_upgrade_protocol_with_inner_call_do_not_modify(self, tmpdir):
input_code = """import ssl
ssl.SSLContext(foo(ssl.PROTOCOL_SSLv2))
"""
expected_output = input_code
self.run_and_assert(tmpdir, input_code, expected_output)
def test_upgrade_protocol_in_expression_do_not_modify(self, tmpdir):
input_code = """import ssl
ssl.SSLContext(ssl.PROTOCOL_SSLv2 if condition else ssl.PROTOCOL_TLS_CLIENT)
"""
expected_output = input_code
self.run_and_assert(tmpdir, input_code, expected_output)
def test_import_no_protocol(self, tmpdir):
input_code = """import ssl
context = ssl.SSLContext()
"""
expected_output = """import ssl
context = ssl.SSLContext(protocol=ssl.PROTOCOL_TLS_CLIENT)
"""
self.run_and_assert(tmpdir, input_code, expected_output)
def test_from_import_no_protocol(self, tmpdir):
input_code = """from ssl import SSLContext
SSLContext()
"""
expected_output = """from ssl import SSLContext
import ssl
SSLContext(protocol=ssl.PROTOCOL_TLS_CLIENT)
"""
self.run_and_assert(tmpdir, input_code, expected_output)