/
websnake.py
208 lines (160 loc) · 6.6 KB
/
websnake.py
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
from untwisted.iostd import LOAD, CLOSE, CONNECT, CONNECT_ERR, \
Client, Stdin, Stdout, lose, create_client
from urllib.parse import urlencode, urlparse
from untwisted.iossl import SSL_CONNECT, create_client_ssl
from untwisted.splits import AccUntil, TmpFile
from untwisted.network import Spin, xmap, spawn, SSL
from untwisted.dispatcher import Dispatcher
from untwisted.event import get_event
from base64 import encodebytes
from untwisted import core
from tempfile import TemporaryFile
from socket import getservbyname
import sys
class Headers(dict):
def __init__(self, data):
for ind in data:
field, sep, value = ind.partition(':')
self[field.lower()] = value
class TransferHandle(object):
DONE = get_event()
def __init__(self, spin):
xmap(spin, AccUntil.DONE, lambda spin, response, data:
spawn(spin, TransferHandle.DONE, Response(response), data))
class ResponseHandle(object):
DONE = get_event()
MAX_SIZE = 1024 * 1024
def __init__(self, spin):
self.response = None
xmap(spin, TransferHandle.DONE, self.process)
def process(self, spin, response, data):
self.response = response
# These handles have to be mapped here
# otherwise it may occur of TmpFile spawning
# done in the first cycle.
xmap(spin, TmpFile.DONE, lambda spin, fd, data: lose(spin))
xmap(spin, TmpFile.DONE,
lambda spin, fd, data: fd.seek(0))
xmap(spin, TmpFile.DONE, lambda spin, fd, data:
spawn(spin, ResponseHandle.DONE, self.response))
TmpFile(spin, data, int(response.headers.get(
b'content-length', self.MAX_SIZE)), response.fd)
# Reset the fd in case of the socket getting closed
# and there is no content-length header.
xmap(spin, CLOSE, lambda spin, err:
self.response.fd.seek(0))
# The fact of destroying the Spin instance when on TmpFile.DONE
# doesnt warrant the CLOSE event will not be fired.
# So it spawns ResponseHandle.DONE only if there is
# a content-length header.
if not b'content-length' in response.headers:
xmap(spin, CLOSE, lambda spin, err:
spawn(spin, ResponseHandle.DONE, self.response))
class Response(object):
def __init__(self, data):
data = data.decode('utf8')
data = data.split('\r\n')
response = data.pop(0)
self.version, self.code, self.reason = response.split(' ', 2)
self.headers = Headers(data)
self.fd = TemporaryFile('w+b')
class HttpCode(object):
def __init__(self, spin):
xmap(spin, ResponseHandle.DONE, lambda spin, response:
spawn(spin, response.code, response))
def on_connect(spin, request):
AccUntil(spin)
TransferHandle(spin)
ResponseHandle(spin)
HttpCode(spin)
spin.dump(request)
def create_con_ssl(addr, port, data):
con = create_client_ssl(addr, port)
xmap(con, SSL_CONNECT, on_connect, data)
return con
def create_con(addr, port, data):
con = create_client(addr, port)
xmap(con, CONNECT, on_connect, data)
return con
def build_headers(headers):
data = ''
for key, value in headers.items():
data = data + '%s: %s\r\n' % (key, value)
data = data + '\r\n'
return data
class Context(Dispatcher):
def __init__(self, method, addr, *args, **kwargs):
self.addr = addr
self.args = args
self.kwargs = kwargs
self.method = method
self.con = method(addr, *self.args, **self.kwargs)
self.con.add_map(ResponseHandle.DONE, self.redirect)
super(Context, self).__init__()
self.con.add_handle(self.proxy)
def redirect(self, con, response):
location = response.headers.get('location')
if location: self.switch(location)
def switch(self, location):
con = self.method(location, *self.args, **self.kwargs)
con.add_map(ResponseHandle.DONE, self.redirect)
con.add_handle(self.proxy)
def proxy(self, con, event, args):
self.drive(event, con, *args)
class ContextGet(Context):
def __init__(self, addr, args={},
headers={}, version='HTTP/1.1', auth=()):
super(ContextGet, self).__init__(get, addr,
args, headers, version, auth)
class ContextPost(Context):
def __init__(self, addr, payload='',
version='HTTP/1.1', headers={}, auth=()):
super(ContextPost, self).__init__(post, addr,
payload, version, headers, auth)
def get(addr, args={}, headers={}, version='HTTP/1.1', auth=()):
"""
It does an http/https request.
"""
addr = addr.strip().rstrip()
url = urlparse(addr)
default = {
'user-agent':'Websnake/1.0.0',
'accept-charset': 'ISO-8859-1,utf-8;q=0.7,*;q=0.7',
'connection':'close',
'host': url.hostname}
default.update(headers)
args = '?%s' % urlencode(args) if args else ''
if auth: default['authorization'] = build_auth(*auth)
data = 'GET %s%s %s\r\n' % (url.path + ('?' + url.query if \
url.query else ''), args, version)
data = data + build_headers(default)
port = url.port if url.port else getservbyname(url.scheme)
data = data.encode('ascii')
return create_con_ssl(url.hostname, port, data) \
if url.scheme == 'https' else create_con(url.hostname, port, data)
def post(addr, payload=b'', version='HTTP/1.1', headers={}, auth=()):
"""
"""
addr = addr.strip().rstrip()
url = urlparse(addr)
default = {'user-agent':"Untwisted-requests/1.0.0",
'accept-charset':b'ISO-8859-1,utf-8;q=0.7,*;q=0.7',
'connection':'close',
'host': url.hostname,
'content-type': 'application/x-www-form-urlencoded',
'content-length': len(payload)}
default.update(headers)
request = 'POST %s %s\r\n' % (url.path + ('?' + url.query if \
url.query else ''), version)
if auth: default['authorization'] = build_auth(*auth)
request = (request + build_headers(default)).encode('ascii') + payload
port = url.port if url.port else getservbyname(url.scheme)
return create_con_ssl(url.hostname, port, request) \
if url.scheme == 'https' else create_con(url.hostname, port, request)
def build_auth(username, password):
# The headers will be encoded as utf8.
username = username.encode('utf8')
password = password.encode('utf8')
base = encodebytes(b'%s:%s' % (username, password))
base = base.replace(b'\n', b'').decode('utf8')
return "Basic %s" % base