annotator.py

import csv
import json
import sys
import re
import en_core_web_lg


def usage():
    print 'usage: python annotator.py <path_to_dataset_in_csv> <output_file>'


def _get_annotations(desc, pattern, label):
    regex = re.compile(pattern)
    return [(match.start(), match.start() + len(match.group()), label) for match in regex.finditer(desc)]


# ------------------- ANNOTATION RULES -------------------

def find_functions(desc):
    pattern = r'[a-zA-Z0-9]+(_[a-zA-Z0-9]+)+'
    return _get_annotations(desc, pattern, "FUNCTION")


def find_kernel_version(desc):
    pattern = r'\d+(\.\d+)+(-?)(\w+)?'
    return _get_annotations(desc, pattern, "VERSION")


def find_file_path(desc):
    pattern = r'(/)?\w+(/\w+)+\.c'
    return _get_annotations(desc, pattern, "SOURCECODE")


def find_fd_driver(desc):
    pattern = r'(/)?\w+(/\w+(?!\.c))+'
    return _get_annotations(desc, pattern, "DRIVER")


def find_driver(desc):
    pattern = r'[a-zA-Z]+\s[a-zA-Z\-]+\s(?=driver)'
    return _get_annotations(desc, pattern, "DRIVER")


def find_structs(desc):
    pattern = r'(?<=struct\s)\w+(\s\w+)'
    return _get_annotations(desc, pattern, "STRUCT")

# ------------------- END ANNOTATION RULES -------------------


TAGGERS = [
    find_functions,
    find_kernel_version,
    find_file_path,
    find_driver,
    find_structs,
    find_fd_driver
]


def annotate_NER(dataset_file):
    TRAIN_DATA = []
    with open(dataset_file, 'r') as cve_dataset_f:
        # cve_reader = csv.DictReader(cve_dataset_f, delimiter=';')
        for cve in cve_dataset_f:
            entities = []
            for tagger in TAGGERS:
                entities += tagger(cve)
            TRAIN_DATA += [[cve, {'entities': entities}]]
    with open('annotated_{}_NER_train.json'.format(dataset_file.replace('.csv', '')), 'w') as annotated_f:
        json.dump(TRAIN_DATA, annotated_f)


def annotate_DEP(dataset_file):
    TRAIN_DATA = []
    model = en_core_web_lg.load()
    with open(dataset_file, 'r') as cve_dataset_f:
        cve_reader = csv.DictReader(cve_dataset_f, delimiter=';')
        for cve in cve_reader:
            tagged_desc = model(unicode(cve['Avail.']))
            heads = [tok.head.i for tok in tagged_desc]
            deps = ['-']*len(heads)
            TRAIN_DATA += [[cve['Avail.'], {'heads': heads, 'deps': deps}]]
    with open('annotated_{}_DEP_train.json'.format(dataset_file.replace('.csv', '')), 'w') as annotated_f:
        json.dump(TRAIN_DATA, annotated_f)


def main(dataset_file):
    annotate_NER(dataset_file)
    # annotate_DEP(dataset_file)


if __name__ == '__main__':
    if len(sys.argv) <= 1:
        usage()
        sys.exit(0)
    main(sys.argv[1])