upload/export.py - master - socialgraph/dataset

upload/export.py
121 строка4 KB
socialgraph
parse
09 авг 2026, 10:22
09 авг 2026, 10:2276cada1
import json
import os
import re

from prepare.utils import write_examples
from tools.const import WORK_DIR


def run(datasets, years):
    work_dir = f'{WORK_DIR}/src'

    examples = []
    for dataset in datasets:
        folder = f'{work_dir}/{dataset}'
        for file in os.listdir(folder):
            if file.endswith('.json'):
                with open(f'{folder}/{file}', encoding='utf-8') as src:
                    src = json.load(src)
                    if src['published'][:4] not in years:
                        continue

                    if 'withdrawn' in src:
                        continue

                    aliases = [q for q in src.get('aliases', []) if q.startswith('CVE')]
                    related = [q for q in src.get('related', []) if q.startswith('CVE')]

                    # if aliases and related and aliases != related:
                    #     print('aliases != related:', aliases, related)

                    if 'references' not in src:
                        continue

                    fixes = {}
                    for q in src['references']:
                        if q['type'] == 'PACKAGE':
                            m = re.search(r'https://github.com/([-\w]+/[-\w]+)', q['url'])
                            if not m:
                                m = re.search(r'([-\w]+/[-\w]+)', q['url'])

                            if m:
                                fixes[m[1]] = []
                            # else:
                            #     print('PACKAGE:', q['url'])

                    if len(fixes) > 1:
                        print('packages:', ', '.join(fixes.keys()))

                    if not fixes:
                        for q in src['references']:
                            if q['type'] == 'WEB':
                                m = re.search(r'https://github.com/([-\w]+/[-\w]+)/commit/\w+', q['url'])
                                if m:
                                    fixes[m[1]] = []

                        # if len(fixes) > 1:
                        #     print('projects:', ', '.join(fixes.keys()))

                    advisory = []
                    for q in src['references']:
                        if q['type'] == 'ADVISORY':
                            m = re.search(r'https://nvd.nist.gov/vuln/detail/(CVE-\d+-\d+)', q['url'])
                            if m:
                                # if cve_ids and m[1] not in cve_ids:
                                #     print('ADVISORY:', m[1])
                                advisory.append(m[1])
                        elif q['type'] == 'WEB':
                            for project in fixes.keys():
                                m = re.search(rf'https://github.com/{project}/commit/(\w+)', q['url'])
                                if m:
                                    fixes[project].append(m[1])

                    # if not cve_ids:
                    #     print(src)

                    cve_ids = aliases or related or advisory

                    for project, commits in fixes.items():
                        if not commits:
                            continue
                        dst = {
                            'id': src['id'],
                            'project': project,
                            'published': src['published'],
                            'cve_ids': cve_ids,
                            'details': src['details'],
                            'commits': commits
                        }

                        if 'summary' in src:
                            dst['summary'] = src['summary']

                        if 'database_specific' in src:
                            if 'cwe_ids' in src['database_specific']:
                                dst['cwe_ids'] = src['database_specific']['cwe_ids']
                            # else:
                            #     print('database_specific:', src['database_specific'])

                        examples.append(dst)

    write_examples(examples, work_dir)


if __name__ == '__main__':
    import argparse

    args = argparse.ArgumentParser()

    args.add_argument(
        "--datasets",
        type=str,
        default="Maven,PyPI"
    )
    args.add_argument(
        "--years",
        type=str,
        default="2026"
    )

    args = args.parse_args()
    run(args.datasets.split(','), args.years.split(','))