conjurer/conjurer_librarian/scrape_bot.py

"""
This module contains the code for the scrape bot.
"""

import json
import logging
import os
import random
import re
import time
from json import JSONDecodeError
from queue import Queue
from threading import Thread
from urllib.request import urlopen

from requests import ConnectionError as RequestsConnectionError
from requests import ConnectTimeout, Timeout

SCR_DATABASE_PATH = os.getenv("CONJURER_LIBRARIAN_DB_PATH", r"C:\\Database\\chunks\\")
SCR_FILENAME = os.getenv("CONJURER_LIBRARIAN_SCRAPE_CHUNK", "40_chunk.txt")
SCR_ENCODING = os.getenv("CONJURER_ENCODING", "utf-8")

WORK_Q = Queue()
random.seed()


def load_ndb_to_q(logger):
    """
    This function loads items from the not_in_db.json file into the work queue.
    """
    logger.info("Loader started")
    while True:
        with open("not_in_db.json", "r+", encoding="utf-8") as ndb_file:
            try:
                ndb_database = json.load(ndb_file)
                for _ in range (1,10):
                    try:
                        key = next(iter(ndb_database))
                        _ = ndb_database.pop(key)
                        logger.info(key)
                        url = f"https://sci-hub.se/{key}"
                        WORK_Q.put([key, url, False])
                    except StopIteration:
                        break
                ndb_file.truncate(0)
                ndb_file.seek(0)
                json.dump(ndb_database, ndb_file, indent=4)
            except JSONDecodeError:
                time.sleep(60 * 60 * 3)
        time.sleep(60*60*3)


def check_if_exists_brute_force(logger):
    """
    This function checks if a document exists using brute force search.
    """
    # Function code here
    while True:
        logger.info("Scraper tick")
        item = WORK_Q.get()
        page_url = item[1]
        logger.error("REFINE: Brute force search!")
        if not page_url.startswith(("http:", "https:")):
            raise ValueError("URL must start with 'http:' or 'https:'")
        blocked = True
        try:
            # trunk-ignore(bandit/B310)
            with urlopen(page_url) as response:
                data = response.read()
                text = data.decode("utf-8")
                if item[2]:
                    logger.info("Already found")
                    time.sleep(180)
                    blocked = False
                else:
                    for line in text.splitlines():
                        if re.match(
                            r".*Unfortunately, Sci-Hub doesn't have the requested document.*",
                            line,
                        ):
                            blocked = False
                            logger.info("Not found")
                            item[2] = False
                        if m := re.match(
                            r".*<embed type=\"application.pdf\"\s*src=\"(.*\.pdf)", line
                        ):
                            blocked = False
                            direct_download_link = "https:" + str(m.group(1))
                            logger.info(direct_download_link)
                            with open(
                                SCR_DATABASE_PATH + SCR_FILENAME,
                                "a",
                                encoding=SCR_ENCODING,
                            ) as operated_file:
                                operated_file.write("\n")
                                operated_file.write(item[0])
                                item[2] = True
        except (
            Timeout,
            ConnectTimeout,
            ConnectionRefusedError,
            ConnectionError,
            RequestsConnectionError,
        ):
            pass
        if blocked:
            logger.info(item)
            logger.error("Got blocked. Fuck.")
            time.sleep(60 * 60)
        # trunk-ignore(bandit/B311)
        rand = random.randint(1, 60)
        logger.info(f"Sleeping for {2*rand} minutes")
        time.sleep(120 * rand)


def scraper(logger=None):
    """
    This function is responsible for scraping data.
    """
    if not logger:
        logger = logging.getLogger()
        logger.setLevel("DEBUG")
        h1 = logging.StreamHandler()
        logger.addHandler(h1)
        logger.info("No logger set. Test run")

    loader = Thread(target=load_ndb_to_q, args=(logger,))
    _scraper = Thread(target=check_if_exists_brute_force, args=(logger,))
    loader.start()
    _scraper.start()
    loader.join()
    _scraper.join()


if __name__ == "__main__":
    scraper()