]> jfr.im git - yt-dlp.git/blame - youtube_dl/socks.py
[socks] Add socks.py from @bluec0re's public domain implementation
[yt-dlp.git] / youtube_dl / socks.py
CommitLineData
4350b745
YCH
1# This is free and unencumbered software released into the public domain.
2#
3# Anyone is free to copy, modify, publish, use, compile, sell, or
4# distribute this software, either in source code form or as a compiled
5# binary, for any purpose, commercial or non-commercial, and by any
6# means.
7#
8# In jurisdictions that recognize copyright laws, the author or authors
9# of this software dedicate any and all copyright interest in the
10# software to the public domain. We make this dedication for the benefit
11# of the public at large and to the detriment of our heirs and
12# successors. We intend this dedication to be an overt act of
13# relinquishment in perpetuity of all present and future rights to this
14# software under copyright law.
15#
16# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
17# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
18# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
19# IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
20# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
21# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
22# OTHER DEALINGS IN THE SOFTWARE.
23#
24# For more information, please refer to <http://unlicense.org/>
25#
26# Example:
27# import socks
28# import ftplib
29# import socket
30#
31# socks.patch_socket()
32#
33# f = ftplib.FTP('ftp.kernel.org')
34# f.login()
35# print f.retrlines('LIST')
36# f.quit()
37#
38# s = socket.create_connection(('www.google.com', 80))
39# s.sendall('HEAD / HTTP/1.0\r\n\r\n')
40# print s.recv(1024)
41# s.close()
42from __future__ import unicode_literals
43import os
44import struct
45import socket
46import time
47
48__author__ = 'Timo Schmid <coding@timoschmid.de>'
49
50_orig_socket = socket.socket
51
52try:
53 from collections import namedtuple
54except ImportError:
55 from Collections import namedtuple
56
57try:
58 from urllib.parse import urlparse
59except:
60 from urlparse import urlparse
61
62try:
63 from enum import Enum
64except ImportError:
65 Enum = object
66
67
68class ProxyError(IOError): pass
69class Socks4Error(ProxyError):
70 CODES = {
71 0x5B: 'request rejected or failed',
72 0x5C: 'request rejected becasue SOCKS server cannot connect to identd on the client',
73 0x5D: 'request rejected because the client program and identd report different user-ids'
74 }
75 def __init__(self, code=None, msg=None):
76 if code is not None and msg is None:
77 msg = self.CODES.get(code)
78 if msg is None:
79 msg = 'unknown error'
80 super(Socks4Error, self).__init__(code, msg)
81
82class Socks5Error(Socks4Error):
83 CODES = {
84 0x01: 'general SOCKS server failure',
85 0x02: 'connection not allowed by ruleset',
86 0x03: 'Network unreachable',
87 0x04: 'Host unreachable',
88 0x05: 'Connection refused',
89 0x06: 'TTL expired',
90 0x07: 'Command not supported',
91 0x08: 'Address type not supported',
92 0xFE: 'unknown username or invalid password',
93 0xFF: 'all offered authentication methods were rejected'
94 }
95
96class ProxyType(Enum):
97 SOCKS4 = 0
98 SOCKS4A = 1
99 SOCKS5 = 2
100
101Proxy = namedtuple('Proxy', ('type', 'host', 'port', 'username', 'password', 'remote_dns'))
102
103_default_proxy = None
104
105def setdefaultproxy(proxytype=None, addr=None, port=None, rdns=True, username=None, password=None, allow_env_override=True):
106 global _default_proxy
107 if allow_env_override:
108 all_proxy = os.environ.get('ALL_PROXY', os.environ.get('all_proxy'))
109 if all_proxy:
110 all_proxy = urlparse(all_proxy)
111 if all_proxy.scheme.startswith('socks'):
112 if all_proxy.scheme == 'socks' or all_proxy.scheme == 'socks4':
113 proxytype = ProxyType.SOCKS4
114 elif all_proxy.scheme == 'socks4a':
115 proxytype = ProxyType.SOCKS4A
116 elif all_proxy.scheme == 'socks5':
117 proxytype = ProxyType.SOCKS5
118 addr = all_proxy.hostname
119 port = all_proxy.port
120 username = all_proxy.username
121 password = all_proxy.password
122
123 if proxytype is not None:
124 _default_proxy = Proxy(proxytype, addr, port, username, password, rdns)
125
126
127def wrap_socket(sock):
128 return socksocket(_sock=sock._sock)
129
130def wrap_module(module):
131 if hasattr(module, 'socket'):
132 sock = module.socket
133 if isinstance(sock, socket.socket):
134 module.socket = sockssocket
135 elif hasattr(socket, 'socket'):
136 socket.socket = sockssocket
137
138def patch_socket():
139 import sys
140 if 'socket' not in sys.modules:
141 import socket
142 sys.modules['socket'].socket = sockssocket
143
144
145class sockssocket(socket.socket):
146 def __init__(self, *args, **kwargs):
147 self.__proxy = None
148 if 'proxy' in kwargs:
149 self.__proxy = kwargs['proxy']
150 del kwargs['proxy']
151 super(sockssocket, self).__init__(*args, **kwargs)
152
153 @property
154 def _proxy(self):
155 if self.__proxy:
156 return self.__proxy
157 return _default_proxy
158
159 @property
160 def _proxy_port(self):
161 if self._proxy:
162 if self._proxy.port:
163 return self._proxy.port
164 return 1080
165 return None
166
167 def setproxy(self, proxytype=None, addr=None, port=None, rdns=True, username=None, password=None):
168 if proxytype is None:
169 self.__proxy = None
170 else:
171 self.__proxy = Proxy(proxytype, addr, port, username, password, rdns)
172
173 def recvall(self, cnt):
174 data = b''
175 while len(data) < cnt:
176 cur = self.recv(cnt - len(data))
177 if not cur:
178 raise IOError("{0} bytes missing".format(cnt-len(data)))
179 data += cur
180 return data
181
182 def _setup_socks4(self, address, is_4a=False):
183 destaddr, port = address
184
185 try:
186 ipaddr = socket.inet_aton(destaddr)
187 except socket.error:
188 if is_4a and self._proxy.remote_dns:
189 ipaddr = struct.pack('!BBBB', 0, 0, 0, 0xFF)
190 else:
191 ipaddr = socket.inet_aton(socket.gethostbyname(destaddr))
192
193 packet = struct.pack('!BBH', 0x4, 0x1, port) + ipaddr
194 if self._proxy.username:
195 username = self._proxy.username
196 if hasattr(username, 'encode'):
197 username = username.encode()
198 packet += struct.pack('!{0}s'.format(len(username)+1), username)
199 else:
200 packet += b'\x00'
201
202 if is_4a and self._proxy.remote_dns:
203 if hasattr(destaddr, 'encode'):
204 destaddr = destaddr.encode()
205 packet += struct.pack('!{0}s'.format(len(destaddr)+1), destaddr)
206
207 self.sendall(packet)
208
209 packet = self.recvall(8)
210 nbyte, resp_code, dstport, dsthost = struct.unpack('!BBHI', packet)
211
212 # check valid response
213 if nbyte != 0x00:
214 self.close()
215 raise ProxyError(0, "Invalid response from server. Expected {0:02x} got {1:02x}".format(0, nbyte))
216
217 # access granted
218 if resp_code != 0x5a:
219 self.close()
220 raise Socks4Error(resp_code)
221
222 def _setup_socks5(self, address):
223 destaddr, port = address
224
225 try:
226 ipaddr = socket.inet_aton(destaddr)
227 except socket.error:
228 if self._proxy.remote_dns:
229 ipaddr = None
230 else:
231 ipaddr = socket.inet_aton(socket.gethostbyname(destaddr))
232
233 auth_methods = 1
234 if self._proxy.username and self._proxy.password:
235 # two auth methods available
236 auth_methods = 2
237 packet = struct.pack('!BBB', 0x5, auth_methods, 0x00) # no auth
238 if self._proxy.username and self._proxy.password:
239 packet += struct.pack('!B', 0x02) # user/pass auth
240
241 self.sendall(packet)
242
243 packet = self.recvall(2)
244 version, method = struct.unpack('!BB', packet)
245
246 # check valid response
247 if version != 0x05:
248 self.close()
249 raise ProxyError(0, "Invalid response from server. Expected {0:02x} got {1:02x}".format(5, version))
250
251 # no auth methods
252 if method == 0xFF:
253 self.close()
254 raise Socks5Error(method)
255
256 # user/pass auth
257 if method == 0x01:
258 username = self._proxy.username
259 if hasattr(username, 'encode'):
260 username = username.encode()
261 password = self._proxy.password
262 if hasattr(password, 'encode'):
263 password = password.encode()
264 packet = struct.pack('!BB', 1, len(username)) + username
265 packet += struct.pack('!B', len(password)) + password
266 self.sendall(packet)
267
268 packet = self.recvall(2)
269 version, status = struct.unpack('!BB', packet)
270
271 if version != 0x01:
272 self.close()
273 raise ProxyError(0, "Invalid response from server. Expected {0:02x} got {1:02x}".format(1, version))
274
275 if status != 0x00:
276 self.close()
277 raise Socks5Error(1)
278 elif method == 0x00: # no auth
279 pass
280
281
282 packet = struct.pack('!BBB', 5, 1, 0)
283 if ipaddr is None:
284 if hasattr(destaddr, 'encode'):
285 destaddr = destaddr.encode()
286 packet += struct.pack('!BB', 3, len(destaddr)) + destaddr
287 else:
288 packet += struct.pack('!B', 1) + ipaddr
289 packet += struct.pack('!H', port)
290
291 self.sendall(packet)
292
293 packet = self.recvall(4)
294 version, status, _, atype = struct.unpack('!BBBB', packet)
295
296 if version != 0x05:
297 self.close()
298 raise ProxyError(0, "Invalid response from server. Expected {0:02x} got {1:02x}".format(5, version))
299
300 if status != 0x00:
301 self.close()
302 raise Socks5Error(status)
303
304 if atype == 0x01:
305 destaddr = self.recvall(4)
306 elif atype == 0x03:
307 alen = struct.unpack('!B', self.recv(1))[0]
308 destaddr = self.recvall(alen)
309 elif atype == 0x04:
310 destaddr = self.recvall(16)
311 destport = struct.unpack('!H', self.recvall(2))[0]
312
313 def _make_proxy(self, connect_func, address):
314 if self._proxy.type == ProxyType.SOCKS4:
315 result = connect_func(self, (self._proxy.host, self._proxy_port))
316 if result != 0 and result is not None:
317 return result
318 self._setup_socks4(address)
319 elif self._proxy.type == ProxyType.SOCKS4A:
320 result = connect_func(self, (self._proxy.host, self._proxy_port))
321 if result != 0 and result is not None:
322 return result
323 self._setup_socks4(address, is_4a=True)
324 elif self._proxy.type == ProxyType.SOCKS5:
325 result = connect_func(self, (self._proxy.host, self._proxy_port))
326 if result != 0 and result is not None:
327 return result
328 self._setup_socks5(address)
329 else:
330 return connect_func(self, address)
331
332 def connect(self, address):
333 self._make_proxy(_orig_socket.connect, address)
334
335 def connect_ex(self, address):
336 return self._make_proxy(_orig_socket.connect_ex, address)