diff --git a/hexlib/env.py b/hexlib/env.py index dacdce3..b2c23ab 100644 --- a/hexlib/env.py +++ b/hexlib/env.py @@ -1,5 +1,7 @@ import redis import os +from fake_useragent import UserAgent + from hexlib.log import stdout_logger from hexlib.web import Web @@ -13,10 +15,13 @@ def get_redis(): def get_web(): + ua = UserAgent() + return Web( proxy=os.environ.get("PROXY", None), rps=os.environ.get("RPS", 1), logger=stdout_logger, cookie_file=os.environ.get("COOKIE_FILE", None), - retry_codes=set(int(x) if x else None for x in os.environ.get("RETRY_CODES", "").split(",")) + retry_codes=set(int(x) if x else None for x in os.environ.get("RETRY_CODES", "").split(",")), + ua=ua[os.environ.get("USER_AGENT")] if os.environ.get("USER_AGENT", None) is not None else None ) diff --git a/hexlib/web.py b/hexlib/web.py index 8ac1398..6eca62e 100644 --- a/hexlib/web.py +++ b/hexlib/web.py @@ -111,7 +111,8 @@ def download_file(url, destination, session=None, headers=None, overwrite=False, class Web: - def __init__(self, proxy=None, rps=1, retries=3, logger=None, cookie_file=None, retry_codes=None, session=None): + def __init__(self, proxy=None, rps=1, retries=3, logger=None, cookie_file=None, retry_codes=None, session=None, + ua=None): self._cookie_file = cookie_file self._proxy = proxy self._logger = logger @@ -125,6 +126,9 @@ class Web: self._session = session + if ua is not None: + session.headers["User-Agent"] = ua + if self._cookie_file: self._session.cookies = load_cookiejar(cookie_file) diff --git a/setup.py b/setup.py index 3f1a104..28683ff 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ from setuptools import setup setup( name="hexlib", - version="1.30", + version="1.31", description="Misc utility methods", author="simon987", author_email="me@simon987.net", @@ -13,6 +13,6 @@ setup( ]}, install_requires=[ "ImageHash", "influxdb", "siphash", "python-dateutil", "redis", "orjson", "zstandard", - "u-msgpack-python", "psycopg2-binary" + "u-msgpack-python", "psycopg2-binary", "fake-useragent" ] )