|
16 | 16 | except ImportError: |
17 | 17 | ssl = None |
18 | 18 |
|
19 | | -from unittest import TestCase, skipUnless |
| 19 | +from unittest import mock, TestCase, skipUnless |
20 | 20 | from test import support |
21 | 21 | from test.support import requires_subprocess |
22 | 22 | from test.support import threading_helper |
@@ -1153,50 +1153,30 @@ class TestFtpcpSecurity(TestCase): |
1153 | 1153 | source server's actual peer address instead, the same as FTP.makepasv(). |
1154 | 1154 | """ |
1155 | 1155 |
|
1156 | | - class _FakeSock: |
1157 | | - def __init__(self, peer_host): |
1158 | | - self._peer = (peer_host, 21) |
1159 | | - def getpeername(self): |
1160 | | - return self._peer |
1161 | | - |
1162 | | - class _FakeSource: |
1163 | | - trust_server_pasv_ipv4_address = False |
1164 | | - def __init__(self, advertised_host, real_host): |
1165 | | - self.sock = TestFtpcpSecurity._FakeSock(real_host) |
1166 | | - self._advertised = advertised_host.replace('.', ',') |
1167 | | - def voidcmd(self, cmd): |
1168 | | - pass |
1169 | | - def sendcmd(self, cmd): |
1170 | | - if cmd == 'PASV': |
1171 | | - return '227 Entering Passive Mode (%s,1,2).' % self._advertised |
1172 | | - return '150 ok' |
1173 | | - def voidresp(self): |
1174 | | - pass |
1175 | | - |
1176 | | - class _FakeTarget: |
1177 | | - def __init__(self): |
1178 | | - self.sendport_args = None |
1179 | | - def voidcmd(self, cmd): |
1180 | | - pass |
1181 | | - def sendport(self, host, port): |
1182 | | - self.sendport_args = (host, port) |
1183 | | - def sendcmd(self, cmd): |
1184 | | - return '150 ok' |
1185 | | - def voidresp(self): |
1186 | | - pass |
| 1156 | + def _make_pair(self, *, advertised_host, real_host, trust=False): |
| 1157 | + source = mock.Mock(spec=ftplib.FTP) |
| 1158 | + source.trust_server_pasv_ipv4_address = trust |
| 1159 | + source.sock.getpeername.return_value = (real_host, 21) |
| 1160 | + # PASV replies give the host as comma-separated octets, not dotted. |
| 1161 | + advertised = advertised_host.replace('.', ',') |
| 1162 | + source.sendcmd.side_effect = lambda cmd: ( |
| 1163 | + f'227 Entering Passive Mode ({advertised},1,2).' |
| 1164 | + if cmd == 'PASV' else '150 ok') |
| 1165 | + target = mock.Mock(spec=ftplib.FTP) |
| 1166 | + target.sendcmd.return_value = '150 ok' |
| 1167 | + return source, target |
1187 | 1168 |
|
1188 | 1169 | def test_ftpcp_ignores_untrusted_pasv_host(self): |
1189 | | - source = self._FakeSource('10.0.0.5', '198.51.100.7') |
1190 | | - target = self._FakeTarget() |
| 1170 | + source, target = self._make_pair(advertised_host='10.0.0.5', |
| 1171 | + real_host='198.51.100.7') |
1191 | 1172 | ftplib.ftpcp(source, 'a', target, 'b') |
1192 | | - self.assertEqual(target.sendport_args, ('198.51.100.7', 258)) |
| 1173 | + target.sendport.assert_called_once_with('198.51.100.7', 258) |
1193 | 1174 |
|
1194 | 1175 | def test_ftpcp_trust_server_pasv_ipv4_address(self): |
1195 | | - source = self._FakeSource('10.0.0.5', '198.51.100.7') |
1196 | | - source.trust_server_pasv_ipv4_address = True |
1197 | | - target = self._FakeTarget() |
| 1176 | + source, target = self._make_pair(advertised_host='10.0.0.5', |
| 1177 | + real_host='198.51.100.7', trust=True) |
1198 | 1178 | ftplib.ftpcp(source, 'a', target, 'b') |
1199 | | - self.assertEqual(target.sendport_args, ('10.0.0.5', 258)) |
| 1179 | + target.sendport.assert_called_once_with('10.0.0.5', 258) |
1200 | 1180 |
|
1201 | 1181 |
|
1202 | 1182 | class MiscTestCase(TestCase): |
|
0 commit comments