aiptx 2.0.7__py3-none-any.whl

aipt_v2/intelligence/searchers/github_searcher.py

@@ -0,0 +1,467 @@
+import os
+import time
+import json
+import datetime
+import dotenv
+from tqdm import tqdm
+import requests
+from git import Repo
+from urllib.parse import quote_plus
+from aipt_v2.utils.searchers.Domain_Filter import repository_filter, code_white_list
+from aipt_v2.utils.searchers.Extension_Filter import for_github_repo_file, for_github_code_file
+from aipt_v2.utils.searchers.util import *
+import utils.searchers.github_config as c
+from scipy.stats import norm
+import csv
+from datetime import datetime as dt
+import shutil
+
+dotenv.load_dotenv()
+
+class GithubSearcher:
+    search_limit_remaining = 30
+    search_limit_reset = 0
+
+    core_limit_remaining = 5000
+    core_limit_reset = 0
+
+    USE_PROXY = False
+    proxies = {
+        'http': 'socks5://127.0.0.1:50532',
+        'https': 'socks5://127.0.0.1:50532'        
+    }
+
+    def __init__(self) -> None:
+        self.session = requests.session()
+        self.token = os.getenv('GITHUB_KEY')
+
+    def convert_to_raw(self, github_blob_url):
+        # replace "blob" with "raw" and return new RAW link
+        github_raw_url = github_blob_url.replace("/github.com/", "/raw.githubusercontent.com/")
+        if "blob/" in github_raw_url:
+            github_raw_url = github_raw_url.replace("blob/", "")
+        return github_raw_url
+
+    def filter_items(self, items, keyword, size_limits, loose_mode = False):
+        ###########################
+        # filter nonrelative repo #
+        ###########################
+        related_items = []
+        # print(items)
+        for item in items:
+            name_temp = str(item.get('name'))
+            description_temp = str(item.get('description'))
+            if 'cve' in keyword.lower():
+                if result_matches_cve(keyword, name_temp) or result_matches_cve(keyword, description_temp):
+                    related_items.append(item)
+
+            if loose_mode:
+                related_items.append(item)
+
+        ########################
+        # filter inferior repo #
+        ########################
+        # filter based on list content
+        filtered_items = [item for item in related_items if not any(filter_string in item['html_url'] for filter_string in repository_filter)]
+        # filter based on scoring
+        extracted_items = []
+        # traverse every item
+        for item in filtered_items:
+            # a new dict to store extracted data
+            extracted_data = {}
+            
+            # extract html_url, empty -> skip this item
+            html_url = item.get('html_url')
+            if not html_url:
+                continue
+            
+            # extract description
+            description = item.get('description')
+            if description is None:
+                description_length = 0
+            else:
+                description_length = len(description)
+            
+            # extract open_issues_count; not exist or not valid -> try to extract open_issues
+            open_issues_count = item.get('open_issues_count')
+            if open_issues_count is None or not isinstance(open_issues_count, int):
+                open_issues = item.get('open_issues')
+                if open_issues is not None and isinstance(open_issues, int):
+                    open_issues_count = open_issues
+                else:
+                    open_issues_count = 0
+
+            # extract stargazers_count; not exist or not valid -> try to extract watchers_count; not exist or not valid -> try to extract watchers
+            # update in original item
+            stars_count = item.get('stargazers_count')
+            if stars_count is None or not isinstance(stars_count, int):
+                watchers_count = item.get('watchers_count')
+                if watchers_count is not None and isinstance(watchers_count, int):
+                    stars_count = watchers_count
+                else:
+                    watchers = item.get('watchers')
+                    if watchers is not None and isinstance(watchers, int):
+                        stars_count = watchers
+                    else:
+                        stars_count = 0
+            item['stars_count'] = stars_count
+            stars_count += 1 # to match up with fork calculation
+
+            # extract forks_count; not exist or not valid -> try to extract forks
+            # update in original item
+            forks_count = item.get('forks_count')
+            if forks_count is None or not isinstance(forks_count, int):
+                forks = item.get('forks')
+                if forks is not None and isinstance(forks, int):
+                    forks_count = forks
+                else:
+                    forks_count = 0
+            item['forks_count'] = forks_count
+            forks_count += 1 # avoid using 0 as divisor
+
+            # extract create_date; not exist or not valid -> set a default value
+            create_date = item.get('created_at')
+            if create_date is None or not isinstance(create_date, str):
+                create_date = "2020-01-01T02:28:41Z"
+            
+            # extract topics; allow to be empty, and transfer it to a string splitted by comma
+            topics_list = item.get('topics', [])
+            topics_str = ','.join(topics_list)
+            
+            # add to new dict
+            extracted_data['html_url'] = html_url
+            extracted_data['description_length'] = description_length
+            extracted_data['open_issues_count'] = open_issues_count
+            extracted_data['topics'] = topics_str.lower()
+            extracted_data['stars_count'] = stars_count
+            extracted_data['forks_count'] = forks_count
+            extracted_data['create_date'] = create_date
+            
+            extracted_items.append(extracted_data)
+
+        for item in extracted_items:
+            if item['description_length'] <= 300:
+                d_score = c.max_confs[0]
+            else:
+                d_score = c.times_0 * norm.pdf(item['description_length'], loc=c.mus[0], scale=c.sigmas[0])
+
+            if item['open_issues_count'] <= 30:
+                i_score = c.max_confs[1]
+            else:
+                i_score = c.times_1 * norm.pdf(item['open_issues_count'], loc=c.mus[1], scale=c.sigmas[1])
+
+            if keyword.strip().lower() in item['topics']:
+                t_score = c.max_confs[2]
+            else:
+                t_score = 0.2 # avoid having a score difference that is too large, as it loses its meaning
+
+            item['conf_score'] = d_score * c.conf_score_weights[0] + i_score * c.conf_score_weights[1] + t_score * c.conf_score_weights[2]
+            
+
+            # transfer to datetime object
+            given_time = dt.fromisoformat(item['create_date'].replace("Z", ""))
+            current_time = dt.now()
+            # get the difference in days
+            days_difference = (current_time - given_time).days
+            lamda = c.times_2 * norm.pdf(days_difference, loc=c.mus[2], scale=c.sigmas[2]) + c.base_line
+            item['efct_score'] = lamda * item['stars_count'] / item['forks_count']
+
+
+        filtered_extracted_items = [item for item in extracted_items if item['conf_score'] >= c.threshold]
+
+        effective_count = len(filtered_extracted_items)
+
+        #####################
+        # filter giant repo #
+        #####################
+        size_filtered_items = [item for item in filtered_items if item['size'] <= c.size_limits[2]]
+
+        #############
+        # sort repo #
+        #############
+        # sort in descending order based on efct_score
+        sorted_filtered_extracted_items = sorted(filtered_extracted_items, key=lambda x: x['efct_score'], reverse=True)
+        # merge results
+        html_index = [item['html_url'] for item in sorted_filtered_extracted_items]
+        merged_filtered_items = []
+        for url in html_index:
+            for item in size_filtered_items:
+                if item['html_url'] == url:
+                    merged_filtered_items.append(item)
+
+        # grouping
+        group1 = [item for item in merged_filtered_items if c.size_limits[0] <= item['size'] <= c.size_limits[1]]
+        group2 = [item for item in merged_filtered_items if item['size'] < c.size_limits[0]]
+        group3 = [item for item in merged_filtered_items if item['size'] > c.size_limits[1]]
+
+        # merge groups
+        grouped_filtered_items = []
+        for item in group1:
+            grouped_filtered_items.append(item)
+        for item in group2:
+            grouped_filtered_items.append(item)
+        for item in group3:
+            grouped_filtered_items.append(item)
+
+        return grouped_filtered_items, effective_count
+
+    def _search_code(self, query_type: str = 'repositories',
+                          query_body: str = '',
+                          qualifiers = '',
+                          page: int = 1,
+                          per_page: int = 30,
+                          sort_method: str = '',
+                          url=None,
+                          keyword: str = '', size_limits: int = [0, 1000000],
+                          loose_mode: bool = False ):
+        timestamp = int(datetime.datetime.now().timestamp())
+        # print("called!")
+        if self.search_limit_remaining == 0 and timestamp < self.search_limit_reset + 3:
+            time.sleep(self.search_limit_reset - timestamp + 5)
+
+        if url is None:
+            # URL encode the query_body
+            query_body_encoded = quote_plus(query_body)
+            query_body_encoded = quote_plus(query_body).rstrip('+')
+            if loose_mode: # to get broader match result
+                query_body_encoded = query_body.replace(' ', '+')
+                query_body_encoded = query_body_encoded.replace('"', '')
+            # construct the URL, ensuring that the parameters are separated by &, and check if the parameters exist.
+            url = f'https://api.github.com/search/{query_type}?q={query_body_encoded}'
+            if qualifiers:
+                qualifiers_body_encoded = quote_plus(qualifiers)
+                qualifiers_body_encoded = qualifiers_body_encoded.replace('+', '&')
+                url += f'&{qualifiers_body_encoded}'
+            if sort_method:
+                url += f'&sort={sort_method}&order=desc'
+            if page > 1:
+                url += f'&page={page}'
+            if per_page > 0:
+                url += f'&per_page={per_page}'
+            # print(url)
+
+        header = {
+            'Accept': 'application/vnd.github+json',
+            'Authorization': f'Bearer {self.token}',
+            'X-GitHub-Api-Version': '2022-11-28'
+        }
+
+        if self.USE_PROXY:
+            resp = self.session.get(url, headers=header, proxies=self.proxies)
+        else:
+            resp = self.session.get(url, headers=header)
+        next_page = None
+        # print(resp.text)
+
+        if 'link' in resp.headers.keys():
+            links = resp.headers['link']
+            links = links.split(',')
+            for link in links:
+                addr, rel = link.split(';')
+                addr = addr.strip()[1:-1]
+
+                if rel.find('next') >= 0:
+                    next_page = addr
+
+        if 'X-RateLimit-Remaining' in resp.headers.keys():
+            self.search_limit_remaining = int(
+                resp.headers['X-RateLimit-Remaining'])
+
+        if 'X-RateLimit-Reset' in resp.headers.keys():
+            self.search_limit_reset = int(resp.headers['X-RateLimit-Reset'])
+
+        if 'Retry-After' in resp.headers.keys():
+            after = int(resp.headers['Retry-After'])
+            time.sleep(after + 3)
+
+        if resp.status_code != 200:
+            # print(resp.text)
+            return None, None
+
+        result = json.loads(resp.text)
+        items = result.get('items', [])
+
+        mode = True if loose_mode else False
+        if query_type == 'repositories':
+            filtered_items, effective_count = self.filter_items(items, keyword, size_limits, loose_mode = mode)
+             # update result
+            result['effective_count'] = effective_count
+        elif query_type == 'code':
+            filtered_items = [item for item in items if any(repo in item['html_url'] for repo in code_white_list)]
+             # update result
+            result['effective_count'] = len(filtered_items)
+        # update result
+        result['items'] = filtered_items
+        
+        return result, next_page
+
+    def search_keyword(self, keyword:str, output_dir: str, filter_on: bool = True, loose_mode: bool = False):
+        # print("github called")
+        os.makedirs(output_dir, exist_ok=True)
+        code_result = None
+
+        mode = True if loose_mode else False
+        repo_result, _ = self._search_code(query_type='repositories',
+                                query_body=f'"{keyword}"',
+                                qualifiers='',
+                                per_page=c.per_page,
+                                size_limits=c.size_limits, keyword = keyword, loose_mode = mode)
+        if filter_on:
+            code_result, _ = self._search_code(query_type='code',
+                                    query_body=f'"{keyword}"',
+                                    qualifiers='',
+                                    keyword = keyword, loose_mode = mode)
+        # print(repo_result)
+        # with open ("test.txt", 'w') as f:
+        #     f.write(json.dumps(repo_result, indent=4, ensure_ascii=False))
+        #     f.write(json.dumps(code_result, indent=4, ensure_ascii=False))
+        total_repo = 0
+        total_code = 0
+        if repo_result is None:
+            print("search repo error")
+        else:
+            total_repo = repo_result['effective_count']
+
+        if code_result is None:
+            print("search code error")
+        else:
+            total_code = code_result['effective_count']
+
+        # Maximum files available on Github
+        repo_target_count = min(c.count_each_keyword, total_repo)
+        code_target_count = min(c.count_each_keyword, total_code)
+
+        repo_count = 0 
+        # Repo Download Iteration
+        with tqdm(total=repo_target_count, desc=f'Searching repositories related to {keyword} from GitHub') as pbar:
+            try:
+                for item in repo_result['items']:
+                    repo_name = item['name']
+                    # repo_full_name = item['full_name']
+                    repo_language = item['language']
+                    repo_star = item['stars_count']
+                    repo_url = item['clone_url']
+
+                    repo_directory = os.path.join(output_dir,
+                                                f'{repo_star}_{repo_name}_{repo_language}')
+                    if not os.path.exists(repo_directory):
+                        os.mkdir(repo_directory)
+                        Repo.clone_from(repo_url, repo_directory)
+                        if filter_on:
+                            file_count_0, _ = count_files_and_size(repo_directory)
+                            # filter files, turn off when running pentestagent
+                            remove_files(repo_directory, filter_list=for_github_repo_file, remove_no_extension=True)
+                            file_count_1, _ = count_files_and_size(repo_directory)
+                            if file_count_0 * 0.5 > file_count_1 or file_count_1 > c.file_num_limits:
+                                shutil.rmtree(repo_directory)  # recursively remove directories
+                            else:
+                                repo_count += 1
+                            if repo_count == c.base_limit:
+                                break
+
+
+                    pbar.update()
+
+            except Exception as e:
+                print(e)
+
+        try:
+            repo_result.pop('incomplete_results', None)
+            for item in repo_result['items']:
+                item.update({key: item[key] for key in ['name', 'clone_url', 'stars_count', 'forks_count'] if key in item})
+                for key in list(item.keys()):
+                    if key not in ['name', 'clone_url', 'stars_count', 'forks_count']:
+                        del item[key]
+
+            # repo_result['items'] = [{key: item[key] for key in item if key in ['name', 'clone_url', 'stars_count', 'forks_count']} for item in repo_result['items']]
+            total_stars_count = 0 # initialize
+            total_forks_count = 0 # initialize
+            for item in repo_result['items']:
+                total_stars_count += item['stars_count']
+                total_forks_count += item['forks_count']
+            repo_result['total_stars_count'] = total_stars_count
+            repo_result['total_forks_count'] = total_forks_count
+            if (repo_result['total_count'] <= c.per_page) or (repo_result['total_count'] > c.per_page and repo_result['effective_count'] < c.per_page):
+                repo_result['trend_score'] = c.trend_weights[0] * repo_result['effective_count'] + c.trend_weights[1] * (total_stars_count + total_forks_count)
+            else:
+                repo_result['trend_score'] = c.trend_weights[0] * repo_result['total_count'] \
+                                            + c.trend_weights[1] * (total_stars_count + total_forks_count) * c.alpha
+            try:
+                with open(os.path.join(os.path.dirname(output_dir), "Trend_Score.json"), "w") as f:
+                    f.write(json.dumps(repo_result, indent=4, ensure_ascii=False))
+            except Exception as e:
+                print(e)
+
+        except Exception as e:
+            print(e)
+
+        # Code Download Iteration
+        code_output_dir = os.path.join(output_dir,"Code_File")
+        if not os.path.exists(code_output_dir):
+            os.makedirs(code_output_dir)
+        index_csv_path = os.path.join(code_output_dir, 'index.csv')
+        with tqdm(total=code_target_count, desc=f'Downloading code files related to {keyword} from GitHub') as pbar:
+            # initialize index.csv
+            with open(index_csv_path, 'w', newline='') as csvfile:
+                fieldnames = ['original_path', 'original_name', 'new_name']
+                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+                writer.writeheader()
+            try:
+                for item in code_result['items']:
+                    file_name = item['name']
+                    file_path = item['path']
+                    file_url = self.convert_to_raw(item['html_url'])  # use html_url to download
+                    suffix = file_name.split(".")[-1]
+
+                    # selectively filter files
+                    for i in range(len(for_github_code_file)):
+                        if f".{suffix}" in for_github_code_file[i]:
+                            continue
+
+                    # generate nonrepeatable file name
+                    unique_name = f"{dt.now().strftime('%Y%m%d%H%M%S%f')}.{suffix}"
+
+                    # download file
+                    # Security: Add timeout to prevent indefinite hangs (CWE-400)
+                    response = requests.get(file_url, timeout=30)
+                    if response.status_code == 200:
+                        if not result_matches_cve(keyword, response.content):
+                            continue
+                        # save file content to local
+                        with open(os.path.join(code_output_dir, unique_name), 'wb') as f:
+                            f.write(response.content)
+                        # update index.csv
+                        with open(index_csv_path, 'a', newline='') as csvfile:
+                            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+                            writer.writerow({
+                                'original_path': file_path,
+                                'original_name': file_name,
+                                'new_name': unique_name
+                            })
+                    else:
+                        print(f"Failed to download {file_url}, status code: {response.status_code}")
+
+                    pbar.update()
+
+            except Exception as e:
+                print(e)
+
+            # check if index.csv only has headline
+            with open(index_csv_path, 'r', newline='') as csvfile:
+                reader = csv.reader(csvfile)
+                rows = list(reader)  # extract all lines into a list
+
+            if len(rows) == 1:
+                os.remove(index_csv_path)
+            
+            print("Recursively removing empty directories...")
+            remove_empty_directories(output_dir)
+
+
+# def main():
+#     g = GithubSearcher()
+#     app = "CVE-2024-29847"
+#     g.search_keyword(f"{app} exploit", "/root/try/exp_web_data")
+
+# if __name__ == "__main__":
+#     main()