/pentestgpt/utils/search.py

# This file contains the utility function for performing searches online.
import os

import requests
from dotenv import load_dotenv
from newspaper import Article
from tqdm import tqdm


def parse_url_with_newspaper(url: str) -> str:
    """
    This function parses the content of a URL using the newspaper library.
    :param url: the URL to parse
    :return: the content of the URL, main body only.
    """
    article = Article(url)
    article.download()
    article.parse()
    return article.text


def google_search_keyword(keyword: str, top_n=2) -> list:
    """
    This function performs a Google search using the Google search API.
    :param keyword: the keyword to search
    :return: a list of search results in the format of [(title, link), (title, link), ...]
        return [(None, None)] if no search results are found
    """
    load_dotenv()
    google_search_key = os.getenv("google_search_key")
    google_search_cx = os.getenv("google_search_cx")
    # perform the search
    # example query: GET https://www.googleapis.com/customsearch/v1?key=INSERT_YOUR_API_KEY&cx=017576662512468239146:omuauf_lfve&q=lectures
    url = f"https://www.googleapis.com/customsearch/v1?key={google_search_key}&cx={google_search_cx}&q={keyword}"
    response = requests.get(url).json()
    # get the top n search results on the current page
    try:
        search_results = response["items"][0:top_n]
    except KeyError:  # if no search results are found or the API is out of quota
        return [(None, None)]
    # return the search results
    return [(result["title"], result["link"]) for result in search_results]


def google_search_keyword_openserp(keyword: str, top_n=1) -> list:
    """
    This function performs a Google search via openserp.
    :param keyword: the keyword to search
    :return: a list of search results in the format of [(title, link), (title, link), ...]
        return [(None, None)] if no search results are found
    """
    # maintain a blacklist of websites that should not be included in the search results
    blacklist = ["medium.com", "github.com"]
    # example query: GET http://127.0.0.1:7001/google/search?lang=EN&limit=20&text=hello%20world
    url = f"http://127.0.0.1:7001/google/search?lang=EN&limit=10&text={keyword}"
    try:
        response = requests.get(url, timeout=30).json()  # add timeout of 30 seconds
    except Exception as e:
        print(f"Online Search Request failed: {e};")
        return [(None, None)]
    # get the top n search results on the current page
    search_result = []
    i = 0
    while len(search_result) < top_n and i < len(response):
        try:
            # if any elements in the blacklist are in the link, skip the link
            if any([item in response[i]["url"] for item in blacklist]):
                i += 1
                continue
            search_results = response[i]
            search_result.append((search_results["title"], search_results["url"]))
        except Exception as e:
            print(f"Error: {e}")
        finally:
            i += 1
    return search_result


def crawl_search(search_results: list) -> list:
    """
    This function crawls the search results into a JSON string as RAG.
    :param search_results: the search results returned by `search_online`
        the search result should be in the format of [(title, link), (title, link), ...]
        the search result should be as [None, None] if no search results are found
    :return: a list of strings as RAG
    """
    rag = []
    for title, link in search_results:
        # each website info is in the format of {title: "title", link: "link", content: "content"}
        if title is None or link is None:
            continue

        # try with Jina API read first, then fallback to newspaper, then return failed
        try:
            main_content = jina_read(link)
            if main_content is False:
                main_content = parse_url_with_newspaper(link)
            rag.append({"title": title, "link": link, "content": main_content})
        except Exception as e:
            print(f"Request failed on {link}: {e}")
            rag.append(
                {"title": title, "link": link, "content": "Failed to retrieve content"}
            )
            continue
    return rag


def check_search_connection(backend="google"):
    """
    This function checks if the search backend is available.
    :param backend: the backend to use for searching.
        Default is "openserp".
        Availables are "google" and "openserp"
    """
    if backend == "google":
        # perform a sample search
        try:
            response = google_search_keyword("test")
            return response != [(None, None)]
        except Exception:  # any exception can be handled as False
            return False

        return False
    elif backend == "openserp":
        # perform a get request to localhost:7001
        try:
            response = requests.get(
                "http://localhost:7001/google/search?text=test", timeout=30
            )  ## add timeout of 30 seconds
            return response.status_code == 200
        except Exception:  # any exception can be handled as False
            return False
    else:
        return False


def search_as_RAG(list_of_keywords: list, backend="google") -> list:
    """
    This function searches the list of keywords and returns the search results as RAG.
    :param list_of_keywords: a list of keywords to search
    :param backend: the backend to use for searching.
        Default is "openserp".
        Availables are "google" and "openserp"
    """
    rag = []
    for keyword in tqdm(list_of_keywords):
        if backend == "google":
            search_results = google_search_keyword(keyword)
        elif backend == "openserp":
            search_results = google_search_keyword_openserp(keyword)
        else:
            search_results = google_search_keyword_openserp(keyword)
        rag.extend(crawl_search(search_results))
    return rag


def jina_read(url: str) -> str:
    """
    This function reads the content of a URL using the Jina API.
    :param url: the URL to read
    return
    string format of the URL content
    """

    JINA_API_KEY = os.getenv("JINA_API_KEY")
    base_url = "https://r.jina.ai"
    headers = {
        "Authorization": f"Bearer {JINA_API_KEY}",
    }
    url = f"{base_url}/{url}"
    # make the query
    try:
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            return response.text
        else:
            return False
    except Exception as e:
        return False


def jina_search(text: str) -> str:
    """
    This function performs a search using the Jina API.
    :param url: the URL to search
    return
    string format of the search result
    """
    JINA_API_KEY = os.getenv("JINA_API_KEY")
    base_url = "https://s.jina.ai"
    headers = {
        "Authorization": f"Bearer {JINA_API_KEY}",
    }
    url = f"{base_url}/{text}"
    # make the query
    try:
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            return response.text
        else:
            return False
    except Exception as e:
        return False


if __name__ == "__main__":
    print(jina_search("Common Vulnerabilities in Uniswap V3."))

    # pre-check: check connection
    connection_status = check_search_connection()
    print("Connection Status:", connection_status)
    if connection_status:
        # test 1: search with openserp
        result = google_search_keyword_openserp("AAVE Security Considerations")
        print(result)

        # test 2: crawl information with openserp
        rag = search_as_RAG(
            ["AAVE Security Considerations", "ERC721 Security Considerations"]
        )
        print(rag)
        for item in rag:
            print(item["title"], item["link"])

    else:
        print("Search Backend is not available.")
        # use the Google API search
        rag = search_as_RAG(
            ["AAVE Security Considerations", "ERC721 Security Considerations"],
            backend="google",
        )
        print(rag)