upload/github.py - master - socialgraph/dataset

upload/github.py
73 строки2 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(years):
    work_dir = f'{WORK_DIR}/src'

    examples = []
    folder = f'{work_dir}/GitHub'
    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_at'][:4] not in years:
                    continue

                url = src['source_code_location']
                if not url:
                    continue

                m = re.search(r'https://github.com/([-\w]+/[-\w]+)', url)
                if not m:
                    m = re.search(r'([-\w]+/[-\w]+)', url)
                if not m:
                    continue

                project = m[1]

                commits = []
                for q in src['references']:
                    m = re.search(rf'https://github.com/{project}/commit/(\w+)', q)
                    if m:
                        commits.append(m[1])

                if not commits:
                    continue
                dst = {
                    'id': src['ghsa_id'],
                    'cve_ids': [src['cve_id']] if src['cve_id'] else [],
                    'cwe_ids': [q['cwe_id'] for q in src['cwes']],
                    'project': project,
                    'published': src['published_at'],
                    'details': src['description'],
                    'commits': commits,
                    'severity': src['severity'],
                    'reviewed': src['type'] == 'reviewed'
                }

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

                examples.append(dst)

    write_examples(examples, work_dir)


if __name__ == '__main__':
    import argparse

    args = argparse.ArgumentParser()

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

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