diff --git a/main.py b/main.py index 0055ecc..f606d28 100644 --- a/main.py +++ b/main.py @@ -9,6 +9,7 @@ log.info("Shammer is getting ready...") # Init static vars INDEX_URL = "https://bitscuit.be/" +MAX_CONTENT_LENGTH = 500000 # Create session session = requests.Session() @@ -27,7 +28,16 @@ while not urls.empty(): # Perform request log.info("Fetching '%s'..."%url) try: + r = session.get(url, stream=True) + + # Check file size + if ("Content-Length" in r.headers.keys() and int(r.headers["Content-Length"]) > MAX_CONTENT_LENGTH) or ("content-length" in r.headers.keys() and int(r.headers["content-length"]) > MAX_CONTENT_LENGTH): + log.info("too large response") + continue + + # Download full r = session.get(url) + except Exception as e: log.info("failed") log.info(e)