mirror of
https://github.com/Gallopsled/pwntools.git
synced 2026-10-12 02:02:27 +00:00
Scott Laglerandpeace-maker
5efe909a91
* Add Type Annotations for parts of pwnlib/util * fix: use Callable in pwnlib/util/lists.py * fix: Import Errors * fix: use getLogger in pwnlib/util/net * feat: use collections.abc.Buffer where appropriate * fix: remove collections.abc.Buffer to use Sequence * add CHANGELOG.md entry * Improve packing and fiddling types * Introduce own `BytesLike` type instead of 3.12 Buffer * Improve web.wget types * Use TypedDict in util.net Specify generics properly. * Improve pwnlib.proc types * Improve pwnlib.util.lists types * Fix util.splash types * Make TypeVars private * Improve util.hexdump types * Remove partial types from packing.fit/flat Those require more work to be useful. * Fix implicit wget return * Fix term.text type export * Fix js_escape return type * Fix sh_string doctest types * Fix usage of ssh_channel as ssh_process We access the pid, so a plain ssh_channel is not enough. * Improve util.lists types * Improve packing.dd types * Fix relying on socket.AddressFamily type at runtime The type is only in typeshed and not the real socket module. Don't access it at runtime, only at type checking. * Fix lists.findall with string needle * Fix dd on two files * Update mypy-baseline.txt * Remove testing artifacts from mypy-baseline.txt --------- Co-authored-by: peace-maker <[email protected]>
80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
import os
|
|
import tempfile
|
|
from typing import Any
|
|
|
|
from pwnlib.log import getLogger
|
|
from pwnlib.tubes.buffer import Buffer
|
|
from pwnlib.util.misc import size
|
|
|
|
log = getLogger(__name__)
|
|
|
|
def wget(url: str | bytes, save: str | bytes | bool = False, timeout: float | tuple[float, float] | None = 5.0, **kwargs: Any) -> bytes | None:
|
|
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 (float, tuple): Timeout, in seconds. Can be a single value for both connect and read,
|
|
or a tuple of (connect, read). Set to None to disable timeouts.
|
|
|
|
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 None
|
|
|
|
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, str)):
|
|
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
|