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
|