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
|
# -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import division
import os
import six
import tempfile
from pwnlib.log import getLogger
from pwnlib.tubes.buffer import Buffer
from pwnlib.util.misc import size
log = getLogger(__name__)
def wget(url, save=None, timeout=5, **kwargs):
r"""wget(url, save=None, timeout=5) -> str
Downloads a file via HTTP/HTTPS.
Arguments:
url (str): URL to download
save (str or bool): Name to save as. Any truthy value
will auto-generate a name based on the URL.
timeout (int): Timeout, in seconds
Example:
>>> url = 'https://httpbingo.org/robots.txt'
>>> result = wget(url, timeout=60)
>>> result
b'User-agent: *\nDisallow: /deny\n'
>>> filename = tempfile.mktemp()
>>> result2 = wget(url, filename, timeout=60)
>>> result == open(filename, 'rb').read()
True
"""
import requests
with log.progress("Downloading %r" % url, rate=0.1) as w:
w.status("Making request...")
response = requests.get(url, stream=True, timeout=timeout, **kwargs)
if not response.ok:
w.failure("Got code %s" % response.status_code)
return
total_size = int(response.headers.get('content-length',0))
w.status('0 / %s' % size(total_size))
# Find out the next largest size we can represent as
chunk_size = 1
while chunk_size < (total_size/10):
chunk_size *= 1000
# Count chunks as they're received
buf = Buffer()
# Loop until we have all of the data
for chunk in response.iter_content(chunk_size = 2**10):
buf.add(chunk)
if total_size:
w.status('%s / %s' % (size(buf.size), size(total_size)))
else:
w.status('%s' % size(buf.size))
total_data = buf.get()
# Save to the target file if provided
if save:
if not isinstance(save, (bytes, six.text_type)):
save = os.path.basename(url)
save = save or tempfile.NamedTemporaryFile(dir='.', delete=False).name
with open(save,'wb+') as f:
f.write(total_data)
w.success('Saved %r (%s)' % (f.name, size(total_data)))
else:
w.success('%s' % size(total_data))
return total_data
|
