aiptx 2.0.2__py3-none-any.whl

aipt_v2/intelligence/searchers/google_searcher.py

@@ -0,0 +1,281 @@
+import os
+# from googlesearch import search
+import re
+from tqdm import tqdm
+# import csv
+from deep_translator import GoogleTranslator
+import spacy
+from bs4 import BeautifulSoup
+import requests
+from aipt_v2.utils.searchers.Domain_Filter import domain_filter
+from aipt_v2.utils.searchers.Extension_Filter import for_google_webpage
+from aipt_v2.utils.searchers.util import *
+from langdetect import detect
+import time
+from readability import Document
+from inscriptis import get_text
+import dotenv
+from aipt_v2.utils.model_manager import get_model
+
+dotenv.load_dotenv()
+
+nlp = spacy.load("en_core_web_sm")
+
+def google_search(query, api_key, cse_id, num=10, **kwargs):
+    """
+    use Google Custom Search JSON API toexecute search
+    
+    Args:
+        query (str): search key word
+        api_key (str): your Google API key
+        cse_id (str): your self defined serch engine ID (cx parameter)
+        num (int, optional): num of search results
+        **kwargs: other parameters
+    
+    Returns:
+        dict: a JSON from API
+    """
+    url = "https://www.googleapis.com/customsearch/v1"
+    params = {
+        'q': query,
+        'key': api_key,
+        'cx': cse_id,
+        'num': num
+    }
+    # merge other parameters
+    params.update(kwargs)
+    
+    # Security: Add timeout to prevent indefinite hangs (CWE-400)
+    response = requests.get(url, params=params, timeout=30)
+    response.raise_for_status()  # if request fails
+    return response.json()
+
+class GoogleSearcher:
+    '''Given a query with exploit name and app number, search on Google, get 10 hyperlinks, create folder with 10 different files from 10 links'''
+    def __init__(self, model_name: str = "openai"):
+        self.llm = get_model(model_name)
+
+    def fetch_webpage_content(self,url, timeout=180):
+        try:
+            response = requests.get(url, timeout=timeout) # set timeout management
+            response.raise_for_status()  # check if request is successful
+            return response.text
+        except requests.RequestException as e:
+            print(f"Request Error: {e}")
+            return None
+
+    def translate_text(self, text, source='auto', target='en', max_length=5000):
+        doc = nlp(text)
+        parts = []
+        current_part = ""
+        
+        for sent in doc.sents:
+            if len(current_part) + len(sent.text) <= max_length:
+                current_part += sent.text + " "
+            else:
+                parts.append(current_part)
+                current_part = sent.text + " "
+        if current_part:
+            parts.append(current_part)
+        
+        translated_parts = [GoogleTranslator(source=source, target=target).translate(part.strip()) for part in parts if len(part.strip()) > 0]
+        return ' '.join(translated_parts)
+    
+    def extract_content(self, result, start_marker, end_marker):
+        # define regular expression pattern
+        pattern = f'{start_marker}(.*?){end_marker}'
+        # use re.DOTALL allows . to match any character, including newline characters
+        matches = re.findall(pattern, result, re.DOTALL)
+        return matches
+    
+    def filter_full_file(self, file_content):
+        # use configed LLM to filter files
+        if not self.llm:
+            print("LLM not initialized, skipping file filtering")
+            return "<<@#%_tag useless >>@#%_tag"
+            
+        from langchain_core.messages import HumanMessage, SystemMessage
+        
+        system_message = SystemMessage(content="""You're an excellent cybersecurity penetration tester assistant. 
+        You need to help the tester in filtering a file's content, and your commitment is essential to the task. 
+        You will be provided a file's content, you need to read it, and then use your knowledge to determine if the file is directly useful for an actual penetration.
+        "DIRECTLY USEFUL" means you can directly follow the instructions provided by the file to penetrate.
+        If so, you are supposed to label it, then it will be passed to your co-workers.
+
+        To help you understand what is "DIRECTLY USEFUL", here are some examples: 
+        \t\t"code(s):\n  <?php
+                                // PHP script to test file upload vulnerability
+                                echo "File upload vulnerability test successful";
+                                ?>"
+        \t\t"plain instruction(s):\n   Open the login page of the target website.
+                                Enter the following code in the username or password field: <script>alert('XSS')</script>.
+                                Submit the form.
+                                If there is an XSS vulnerability, a warning box with the word "XSS" will pop up in your browser."
+        \t\t"command line operation(s):\n  curl -d "username=whoami; cat /etc/passwd" http://example.com/login\n
+                                                or\n
+                                                netsh>rpc
+                                                netsh rpc>filter
+                                                netsh rpc filter>add rule layer=um actiontype=block
+                                                netsh rpc filter>add condition field=if_uuid matchtype=equal data=4b324fc8-1670-01d3-1278-5a47bf6ee188
+                                                netsh rpc filter>add filter
+                                                netsh rpc filter>quit"
+        \t\t"payload(s):\n   ' OR 1=1 --    "
+         
+        To make the progress more efficient, we need to predefine a pair of unique strings("<<@#%_tag", ">>@#%_tag") in your answer so that others will recognize them easier.
+        So the format in your answer should looks like: "<<@#%_tag useful >>@#%_tag" or "<<@#%_tag useless >>@#%_tag". The tag should be unique.
+        
+        Remember, you should be careful. A common misleading scenario is a file actually contains something helpful, but it is too long or the information is hidden too deeply that can make it be excluded. So you need to carefully read the whole text. 
+        Meanwhile, to reduce your co-workers' burdens, you need to be strict. It is okay that you find the file is not actually useful to execute penetration. If so, feel free to skip those parts.
+        You do not need to make assumption that a strange URL or link may contain something useful. We can access them through other approaches. So if these things appear in the file, make sure they do not affect your judgement.
+        You can summarize but do not conclude or make assumptions, and your answer should be your most confident one. Keep the answer concise.
+        """)
+
+        user_message = HumanMessage(content=f"""Please make judgement and filter the following content: \n\n{file_content}
+        \n\n\n\n\n
+
+        Please make sure that you have used the unique string pair("<<@#%_tag", ">>@#%_tag"), especially do not forget to add ">>@#%_tag" as an end.""")
+
+        try:
+            response = self.llm.invoke([system_message, user_message])
+            summary = response.content.strip()
+            return summary
+        except Exception as e:
+            print(f"Error calling LLM: {e}")
+            return "<<@#%_tag useless >>@#%_tag"
+
+    def search_keyword(self, keyword: str, output_dir: str):
+        # print("google called")
+        links = []
+        search_keyword = keyword 
+        # domain_filter = ["github", "suibian"]
+        # + " exploit"
+        search_results = []
+        try:
+            time.sleep(5)
+            search_results = google_search(search_keyword, os.environ.get("GOOGLE_API_KEY"), os.environ.get("GOOGLE_CSE_ID"))
+        except:
+            print("Error occurred during google search. Continuing...")
+
+        # search_results = list(search_results)[:10]
+
+        # create list with all links
+        # for result in search_results:
+        #     print(result)
+        #     if all(domain not in result.lower() for domain in domain_filter):
+        #         web_name = result.split("//")[1]
+        #         web_name = web_name.replace("/", "-")
+        #         web_name = web_name[:30] if len(web_name) > 30 else web_name
+        #         links.append((web_name, result))
+        # if not os.path.exists(output_dir):
+
+        if 'items' in search_results:
+            for i, item in enumerate(search_results['items'], 1):
+                
+                # process links - filter with domain name
+                result_link = item['link']
+                if domain_filter:
+                    if all(domain not in result_link.lower() for domain in domain_filter):
+                        web_name = result_link.split("//")[-1].split("/")[0]
+                        web_name = web_name.replace("/", "-")
+                        web_name = web_name[:30] if len(web_name) > 30 else web_name
+                        links.append((web_name, result_link))
+                else:
+                    web_name = result_link.split("//")[-1].split("/")[0]
+                    web_name = web_name.replace("/", "-")
+                    web_name = web_name[:30] if len(web_name) > 30 else web_name
+                    links.append((web_name, result_link))
+        else:
+            print("Cannot find related info.")
+        self.create_directories(output_dir, links)
+    
+    def create_directories(self, output_dir, links):
+        # create a folder if it doesn't exist already
+        # print("called!")
+        os.makedirs(output_dir, exist_ok=True)
+        for (name, link) in tqdm(links, desc="Crawling Google pages"):
+            try:
+                # if exist, then skip
+                if os.path.exists(os.path.join(output_dir, name)):
+                    continue
+
+                time.sleep(5)
+                # document = SimpleWebPageReader(html_to_text=True).load_data([link])[0].text
+                soup = BeautifulSoup(self.fetch_webpage_content(link), "html.parser")
+
+                # extension list
+                extensions_to_remove = for_google_webpage
+
+                # remove images
+                for ule_tag in soup.find_all('True'):
+                    src = ule_tag.get('src')
+                    if src and any(src.lower().endswith(ext) for ext in extensions_to_remove):
+                        ule_tag.decompose()
+
+                # remove links, may not necessary
+                # for a_tag in soup.find_all('a'):
+                #     href = a_tag.get('href')
+                #     if href and any(href.lower().endswith(ext) for ext in extensions_to_remove):
+                #         a_tag.decompose()
+
+                doc = Document(soup.prettify()) # restore to html format, then transfer to the format can be processed by readability
+                content = doc.summary() # summarize, remove nonrelated content, html format
+                full_text = get_text(content) # get clean text, with relative location remained
+
+                # not empty -> creaye directory, else skip
+                if full_text.strip():
+                    os.mkdir(os.path.join(output_dir, name))
+                else:
+                    continue
+
+                # translate Chinese webpages 
+                if full_text.strip():
+                    lan = detect(full_text)
+                    # print(language)
+                    if lan != "en":
+                        time.sleep(5)
+                        full_doc = self.translate_text(full_text)
+                    else:
+                        full_doc = full_text
+
+            except Exception as e:
+                print("Error occurred while downloading web page:", e)
+                continue
+
+            if len(full_doc) > 1000000:
+                continue
+            
+            # use LLM to filter, then create md doc 
+            full_doc_judgement = self.filter_full_file(full_doc)
+            # print(full_doc_judgement)
+            full_doc_judgement = self.extract_content(full_doc_judgement, '<<@#%_tag', '>>@#%_tag')
+            # print(full_doc_judgement)
+            if full_doc_judgement[0].strip() == "useful":
+                with open(os.path.join(output_dir, name, "R_DOC.md"), "w") as f:
+                    f.write(f"link to this page is {link}\n\n")
+                    f.write(full_doc)
+
+        remove_empty_directories(output_dir)
+        
+            
+
+            
+
+# def main():
+#     g = GoogleSearcher()
+#     base_dir = "/root/crawl_classify/data/"
+#     level = "WPNZ"
+#     # with open(os.path.join(base_dir, "index", f"{level}.csv"), mode='r', newline='', encoding='utf-8') as file:
+#     with open(os.path.join(base_dir, "index", "example1.csv"), mode='r', newline='', encoding='utf-8') as file:
+#         reader = csv.reader(file, delimiter='\t')
+
+#         for row in reader:
+#             avd_id, cve_id, zh_keyword, avd_url = row
+#             if cve_id != "N/A":
+#                 g.search_keyword(f"{cve_id} exlpoit", os.path.join(base_dir, level, avd_id, "temp"))
+#             time.sleep(30)
+
+    
+
+# if __name__ == "__main__":
+#     main()
+