Skip to content
This repository was archived by the owner on Apr 14, 2022. It is now read-only.

Commit bc3f5b4

Browse files
RatanShreshthaRatanShreshtha
authored
Merge pull request #119 from pquentin/merge-from-master-2019-06-19
Merge from master (June 6th, 2019)
2 parents 74d6b54 + 36daff0 commit bc3f5b4

8 files changed

+1446
-1624
lines changed

‎test/contrib/test_socks.py

+82-105
Original file line numberDiff line numberDiff line change
@@ -232,13 +232,12 @@ def request_handler(listener):
232232

233233
self._start_server(request_handler)
234234
proxy_url = "socks5://%s:%s" % (self.host, self.port)
235-
pm = socks.SOCKSProxyManager(proxy_url)
236-
self.addCleanup(pm.clear)
237-
response = pm.request("GET", "http://16.17.18.19")
235+
with socks.SOCKSProxyManager(proxy_url) as pm:
236+
response = pm.request("GET", "http://16.17.18.19")
238237

239-
assert response.status == 200
240-
assert response.data == b""
241-
assert response.headers["Server"] == "SocksTestServer"
238+
assert response.status == 200
239+
assert response.data == b""
240+
assert response.headers["Server"] == "SocksTestServer"
242241

243242
def test_local_dns(self):
244243
def request_handler(listener):
@@ -266,13 +265,12 @@ def request_handler(listener):
266265

267266
self._start_server(request_handler)
268267
proxy_url = "socks5://%s:%s" % (self.host, self.port)
269-
pm = socks.SOCKSProxyManager(proxy_url)
270-
self.addCleanup(pm.clear)
271-
response = pm.request("GET", "http://localhost")
268+
with socks.SOCKSProxyManager(proxy_url) as pm:
269+
response = pm.request("GET", "http://localhost")
272270

273-
assert response.status == 200
274-
assert response.data == b""
275-
assert response.headers["Server"] == "SocksTestServer"
271+
assert response.status == 200
272+
assert response.data == b""
273+
assert response.headers["Server"] == "SocksTestServer"
276274

277275
def test_correct_header_line(self):
278276
def request_handler(listener):
@@ -304,10 +302,9 @@ def request_handler(listener):
304302

305303
self._start_server(request_handler)
306304
proxy_url = "socks5h://%s:%s" % (self.host, self.port)
307-
pm = socks.SOCKSProxyManager(proxy_url)
308-
self.addCleanup(pm.clear)
309-
response = pm.request("GET", "http://example.com")
310-
assert response.status == 200
305+
with socks.SOCKSProxyManager(proxy_url) as pm:
306+
response = pm.request("GET", "http://example.com")
307+
assert response.status == 200
311308

312309
def test_connection_timeouts(self):
313310
event = threading.Event()
@@ -317,12 +314,10 @@ def request_handler(listener):
317314

318315
self._start_server(request_handler)
319316
proxy_url = "socks5h://%s:%s" % (self.host, self.port)
320-
pm = socks.SOCKSProxyManager(proxy_url)
321-
self.addCleanup(pm.clear)
322-
323-
with pytest.raises(ConnectTimeoutError):
324-
pm.request("GET", "http://example.com", timeout=0.001, retries=False)
325-
event.set()
317+
with socks.SOCKSProxyManager(proxy_url) as pm:
318+
with pytest.raises(ConnectTimeoutError):
319+
pm.request("GET", "http://example.com", timeout=0.001, retries=False)
320+
event.set()
326321

327322
def test_connection_failure(self):
328323
event = threading.Event()
@@ -333,12 +328,10 @@ def request_handler(listener):
333328

334329
self._start_server(request_handler)
335330
proxy_url = "socks5h://%s:%s" % (self.host, self.port)
336-
pm = socks.SOCKSProxyManager(proxy_url)
337-
self.addCleanup(pm.clear)
338-
339-
event.wait()
340-
with pytest.raises(NewConnectionError):
341-
pm.request("GET", "http://example.com", retries=False)
331+
with socks.SOCKSProxyManager(proxy_url) as pm:
332+
event.wait()
333+
with pytest.raises(NewConnectionError):
334+
pm.request("GET", "http://example.com", retries=False)
342335

343336
def test_proxy_rejection(self):
344337
evt = threading.Event()
@@ -355,12 +348,10 @@ def request_handler(listener):
355348

356349
self._start_server(request_handler)
357350
proxy_url = "socks5h://%s:%s" % (self.host, self.port)
358-
pm = socks.SOCKSProxyManager(proxy_url)
359-
self.addCleanup(pm.clear)
360-
361-
with pytest.raises(NewConnectionError):
362-
pm.request("GET", "http://example.com", retries=False)
363-
evt.set()
351+
with socks.SOCKSProxyManager(proxy_url) as pm:
352+
with pytest.raises(NewConnectionError):
353+
pm.request("GET", "http://example.com", retries=False)
354+
evt.set()
364355

365356
def test_socks_with_password(self):
366357
def request_handler(listener):
@@ -390,14 +381,12 @@ def request_handler(listener):
390381

391382
self._start_server(request_handler)
392383
proxy_url = "socks5://%s:%s" % (self.host, self.port)
393-
pm = socks.SOCKSProxyManager(proxy_url, username="user", password="pass")
394-
self.addCleanup(pm.clear)
395-
396-
response = pm.request("GET", "http://16.17.18.19")
384+
with socks.SOCKSProxyManager(proxy_url, username="user", password="pass") as pm:
385+
response = pm.request("GET", "http://16.17.18.19")
397386

398-
assert response.status == 200
399-
assert response.data == b""
400-
assert response.headers["Server"] == "SocksTestServer"
387+
assert response.status == 200
388+
assert response.data == b""
389+
assert response.headers["Server"] == "SocksTestServer"
401390

402391
def test_socks_with_auth_in_url(self):
403392
"""
@@ -432,14 +421,12 @@ def request_handler(listener):
432421

433422
self._start_server(request_handler)
434423
proxy_url = "socks5://user:pass@%s:%s" % (self.host, self.port)
435-
pm = socks.SOCKSProxyManager(proxy_url)
436-
self.addCleanup(pm.clear)
424+
with socks.SOCKSProxyManager(proxy_url) as pm:
425+
response = pm.request("GET", "http://16.17.18.19")
437426

438-
response = pm.request("GET", "http://16.17.18.19")
439-
440-
assert response.status == 200
441-
assert response.data == b""
442-
assert response.headers["Server"] == "SocksTestServer"
427+
assert response.status == 200
428+
assert response.data == b""
429+
assert response.headers["Server"] == "SocksTestServer"
443430

444431
def test_socks_with_invalid_password(self):
445432
def request_handler(listener):
@@ -452,15 +439,15 @@ def request_handler(listener):
452439

453440
self._start_server(request_handler)
454441
proxy_url = "socks5h://%s:%s" % (self.host, self.port)
455-
pm = socks.SOCKSProxyManager(proxy_url, username="user", password="badpass")
456-
self.addCleanup(pm.clear)
457-
458-
try:
459-
pm.request("GET", "http://example.com", retries=False)
460-
except NewConnectionError as e:
461-
assert "SOCKS5 authentication failed" in str(e)
462-
else:
463-
self.fail("Did not raise")
442+
with socks.SOCKSProxyManager(
443+
proxy_url, username="user", password="badpass"
444+
) as pm:
445+
try:
446+
pm.request("GET", "http://example.com", retries=False)
447+
except NewConnectionError as e:
448+
assert "SOCKS5 authentication failed" in str(e)
449+
else:
450+
self.fail("Did not raise")
464451

465452
def test_source_address_works(self):
466453
expected_port = _get_free_port(self.host)
@@ -492,12 +479,11 @@ def request_handler(listener):
492479

493480
self._start_server(request_handler)
494481
proxy_url = "socks5://%s:%s" % (self.host, self.port)
495-
pm = socks.SOCKSProxyManager(
482+
with socks.SOCKSProxyManager(
496483
proxy_url, source_address=("127.0.0.1", expected_port)
497-
)
498-
self.addCleanup(pm.clear)
499-
response = pm.request("GET", "http://16.17.18.19")
500-
assert response.status == 200
484+
) as pm:
485+
response = pm.request("GET", "http://16.17.18.19")
486+
assert response.status == 200
501487

502488

503489
class TestSOCKS4Proxy(IPV4SocketDummyServerTestCase):
@@ -534,13 +520,12 @@ def request_handler(listener):
534520

535521
self._start_server(request_handler)
536522
proxy_url = "socks4://%s:%s" % (self.host, self.port)
537-
pm = socks.SOCKSProxyManager(proxy_url)
538-
self.addCleanup(pm.clear)
539-
response = pm.request("GET", "http://16.17.18.19")
523+
with socks.SOCKSProxyManager(proxy_url) as pm:
524+
response = pm.request("GET", "http://16.17.18.19")
540525

541-
assert response.status == 200
542-
assert response.headers["Server"] == "SocksTestServer"
543-
assert response.data == b""
526+
assert response.status == 200
527+
assert response.headers["Server"] == "SocksTestServer"
528+
assert response.data == b""
544529

545530
def test_local_dns(self):
546531
def request_handler(listener):
@@ -568,13 +553,12 @@ def request_handler(listener):
568553

569554
self._start_server(request_handler)
570555
proxy_url = "socks4://%s:%s" % (self.host, self.port)
571-
pm = socks.SOCKSProxyManager(proxy_url)
572-
self.addCleanup(pm.clear)
573-
response = pm.request("GET", "http://localhost")
556+
with socks.SOCKSProxyManager(proxy_url) as pm:
557+
response = pm.request("GET", "http://localhost")
574558

575-
assert response.status == 200
576-
assert response.headers["Server"] == "SocksTestServer"
577-
assert response.data == b""
559+
assert response.status == 200
560+
assert response.headers["Server"] == "SocksTestServer"
561+
assert response.data == b""
578562

579563
def test_correct_header_line(self):
580564
def request_handler(listener):
@@ -606,10 +590,9 @@ def request_handler(listener):
606590

607591
self._start_server(request_handler)
608592
proxy_url = "socks4a://%s:%s" % (self.host, self.port)
609-
pm = socks.SOCKSProxyManager(proxy_url)
610-
self.addCleanup(pm.clear)
611-
response = pm.request("GET", "http://example.com")
612-
assert response.status == 200
593+
with socks.SOCKSProxyManager(proxy_url) as pm:
594+
response = pm.request("GET", "http://example.com")
595+
assert response.status == 200
613596

614597
def test_proxy_rejection(self):
615598
evt = threading.Event()
@@ -626,12 +609,10 @@ def request_handler(listener):
626609

627610
self._start_server(request_handler)
628611
proxy_url = "socks4a://%s:%s" % (self.host, self.port)
629-
pm = socks.SOCKSProxyManager(proxy_url)
630-
self.addCleanup(pm.clear)
631-
632-
with pytest.raises(NewConnectionError):
633-
pm.request("GET", "http://example.com", retries=False)
634-
evt.set()
612+
with socks.SOCKSProxyManager(proxy_url) as pm:
613+
with pytest.raises(NewConnectionError):
614+
pm.request("GET", "http://example.com", retries=False)
615+
evt.set()
635616

636617
def test_socks4_with_username(self):
637618
def request_handler(listener):
@@ -659,13 +640,12 @@ def request_handler(listener):
659640

660641
self._start_server(request_handler)
661642
proxy_url = "socks4://%s:%s" % (self.host, self.port)
662-
pm = socks.SOCKSProxyManager(proxy_url, username="user")
663-
self.addCleanup(pm.clear)
664-
response = pm.request("GET", "http://16.17.18.19")
643+
with socks.SOCKSProxyManager(proxy_url, username="user") as pm:
644+
response = pm.request("GET", "http://16.17.18.19")
665645

666-
assert response.status == 200
667-
assert response.data == b""
668-
assert response.headers["Server"] == "SocksTestServer"
646+
assert response.status == 200
647+
assert response.data == b""
648+
assert response.headers["Server"] == "SocksTestServer"
669649

670650
def test_socks_with_invalid_username(self):
671651
def request_handler(listener):
@@ -676,15 +656,13 @@ def request_handler(listener):
676656

677657
self._start_server(request_handler)
678658
proxy_url = "socks4a://%s:%s" % (self.host, self.port)
679-
pm = socks.SOCKSProxyManager(proxy_url, username="baduser")
680-
self.addCleanup(pm.clear)
681-
682-
try:
683-
pm.request("GET", "http://example.com", retries=False)
684-
except NewConnectionError as e:
685-
assert "different user-ids" in str(e)
686-
else:
687-
self.fail("Did not raise")
659+
with socks.SOCKSProxyManager(proxy_url, username="baduser") as pm:
660+
try:
661+
pm.request("GET", "http://example.com", retries=False)
662+
except NewConnectionError as e:
663+
assert "different user-ids" in str(e)
664+
else:
665+
self.fail("Did not raise")
688666

689667

690668
class TestSOCKSWithTLS(IPV4SocketDummyServerTestCase):
@@ -728,10 +706,9 @@ def request_handler(listener):
728706

729707
self._start_server(request_handler)
730708
proxy_url = "socks5h://%s:%s" % (self.host, self.port)
731-
pm = socks.SOCKSProxyManager(proxy_url, ca_certs=DEFAULT_CA)
732-
self.addCleanup(pm.clear)
733-
response = pm.request("GET", "https://localhost")
709+
with socks.SOCKSProxyManager(proxy_url, ca_certs=DEFAULT_CA) as pm:
710+
response = pm.request("GET", "https://localhost")
734711

735-
assert response.status == 200
736-
assert response.data == b""
737-
assert response.headers["Server"] == "SocksTestServer"
712+
assert response.status == 200
713+
assert response.data == b""
714+
assert response.headers["Server"] == "SocksTestServer"

0 commit comments

Comments
 (0)