1919import asyncio
2020import logging
2121import socket
22+ import ipaddress
2223
2324try :
2425 import socks
@@ -37,31 +38,43 @@ class TCP:
3738 TIMEOUT = 10
3839
3940 def __init__ (self , ipv6 : bool , proxy : dict ):
40- self .proxy = proxy
41+ self .socket = None
42+
43+ self .reader = None # type: asyncio.StreamReader
44+ self .writer = None # type: asyncio.StreamWriter
4145
4246 self .lock = asyncio .Lock ()
4347
44- self .socket = socks .socksocket (family = socket .AF_INET6 if ipv6 else socket .AF_INET )
48+ if proxy .get ("enabled" , False ):
49+ hostname = proxy .get ("hostname" , None )
50+ port = proxy .get ("port" , None )
4551
46- self .socket .settimeout (TCP .TIMEOUT )
47-
48- self .reader = None # type: asyncio.StreamReader
49- self .writer = None # type: asyncio.StreamWriter
50- self .proxy_enabled = proxy .get ("enabled" , False )
52+ try :
53+ ip_address = ipaddress .ip_address (hostname )
54+ except ValueError :
55+ self .socket = socks .socksocket (socket .AF_INET )
56+ else :
57+ if isinstance (ip_address , ipaddress .IPv6Address ):
58+ self .socket = socks .socksocket (socket .AF_INET6 )
59+ else :
60+ self .socket = socks .socksocket (socket .AF_INET )
5161
52- if proxy and self .proxy_enabled :
5362 self .socket .set_proxy (
5463 proxy_type = socks .SOCKS5 ,
55- addr = proxy . get ( " hostname" , None ) ,
56- port = proxy . get ( " port" , None ) ,
64+ addr = hostname ,
65+ port = port ,
5766 username = proxy .get ("username" , None ),
5867 password = proxy .get ("password" , None )
5968 )
6069
61- log .info ("Using proxy {}:{}" .format (
62- proxy .get ("hostname" , None ),
63- proxy .get ("port" , None )
64- ))
70+ log .info ("Using proxy {}:{}" .format (hostname , port ))
71+ else :
72+ super ().__init__ (
73+ socket .AF_INET6 if ipv6
74+ else socket .AF_INET
75+ )
76+
77+ self .socket .settimeout (TCP .TIMEOUT )
6578
6679 async def connect (self , address : tuple ):
6780 self .socket .connect (address )
0 commit comments