SoLID/shellcode_i_a32

1# coding=utf-82# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""TODO: Add a description here."""16 17 18import csv19import json20import os21import pandas as pd22import datasets23 24 25# TODO: Add BibTeX citation26# Find for instance the citation on arxiv or on the dataset repo/website27_CITATION = """\28    @inproceedings{liguori-etal-2021-shellcode,29    title = "{S}hellcode{\_}{IA}32: A Dataset for Automatic Shellcode Generation",30    author = "Liguori, Pietro  and31      Al-Hossami, Erfan  and32      Cotroneo, Domenico  and33      Natella, Roberto  and34      Cukic, Bojan  and35      Shaikh, Samira",36    booktitle = "Proceedings of the 1st Workshop on Natural Language Processing for Programming (NLP4Prog 2021)",37    month = aug,38    year = "2021",39    address = "Online",40    publisher = "Association for Computational Linguistics",41    url = "https://aclanthology.org/2021.nlp4prog-1.7",42    doi = "10.18653/v1/2021.nlp4prog-1.7",43    pages = "58--64",44    abstract = "We take the first step to address the task of automatically generating shellcodes, i.e., small pieces of code used as a payload in the exploitation of a software vulnerability, starting from natural language comments. We assemble and release a novel dataset (Shellcode{\_}IA32), consisting of challenging but common assembly instructions with their natural language descriptions. We experiment with standard methods in neural machine translation (NMT) to establish baseline performance levels on this task.",45}46"""47 48# TODO: Add description of the dataset here49# You can copy an official description50_DESCRIPTION = """\51Shellcode_IA32 is a dataset for shellcode generation from English intents. The shellcodes are compilable on Intel Architecture 32-bits.52"""53 54# TODO: Add a link to an official homepage for the dataset here55_HOMEPAGE = "https://github.com/dessertlab/Shellcode_IA32"56 57# TODO: Add the licence for the dataset here if you can find it58_LICENSE = "GNU GENERAL PUBLIC LICENSE"59 60# TODO: Add link to the official dataset URLs here61# The HuggingFace dataset library don't host the datasets but only point to the original files62# This can be an arbitrary nested dict/list of URLs (see below in `_split_generators` method)63_URLs = {64    'default': "https://raw.githubusercontent.com/dessertlab/Shellcode_IA32/main/Shellcode_IA32.tsv",65}66 67 68# TODO: Name of the dataset usually match the script name with CamelCase instead of snake_case69class ShellcodeIA32(datasets.GeneratorBasedBuilder):70    """Shellcode_IA32 a dataset for shellcode generation"""71 72    VERSION = datasets.Version("1.1.0")73 74    # This is an example of a dataset with multiple configurations.75    # If you don't want/need to define several sub-sets in your dataset,76    # just remove the BUILDER_CONFIG_CLASS and the BUILDER_CONFIGS attributes.77 78    # If you need to make complex sub-parts in the datasets with configurable options79    # You can create your own builder configuration class to store attribute, inheriting from datasets.BuilderConfig80    # BUILDER_CONFIG_CLASS = MyBuilderConfig81 82    # You will be able to load one or the other configurations in the following list with83    # data = datasets.load_dataset('my_dataset', 'first_domain')84    # data = datasets.load_dataset('my_dataset', 'second_domain')85    # BUILDER_CONFIGS = [86    #     datasets.BuilderConfig(name="default", version=VERSION, description="This part of my dataset covers the default train/test split"),87    #     #datasets.BuilderConfig(name="second_domain", version=VERSION, description="This part of my dataset covers a second domain"),88    # ]89 90    DEFAULT_CONFIG_NAME = "default"  # It's not mandatory to have a default configuration. Just use one if it make sense.91 92    def _info(self):93        # TODO: This method specifies the datasets.DatasetInfo object which contains informations and typings for the dataset94 95            features = datasets.Features(96                {97                    "intent": datasets.Value("string"),98                    "snippet": datasets.Value("string"),99 100                }101            )102            return datasets.DatasetInfo(103                # This is the description that will appear on the datasets page.104                description=_DESCRIPTION,105                # This defines the different columns of the dataset and their types106                features=features,  # Here we define them above because they are different between the two configurations107                # If there's a common (input, target) tuple from the features,108                # specify them here. They'll be used if as_supervised=True in109                # builder.as_dataset.110                supervised_keys=None,111                # Homepage of the dataset for documentation112                homepage=_HOMEPAGE,113                # License for the dataset if available114                license=_LICENSE,115                # Citation for the dataset116                citation=_CITATION,117            )118 119    def _split_generators(self, dl_manager):120        """Returns SplitGenerators."""121        # TODO: This method is tasked with downloading/extracting the data and defining the splits depending on the configuration122        # If several configurations are possible (listed in BUILDER_CONFIGS), the configuration selected by the user is in self.config.name123 124        # dl_manager is a datasets.download.DownloadManager that can be used to download and extract URLs125        # It can accept any type or nested list/dict and will give back the same structure with the url replaced with path to local files.126        # By default the archives will be extracted and a path to a cached folder where they are extracted is returned instead of the archive127        my_urls = _URLs[self.config.name]128        data_dir = dl_manager.download_and_extract(my_urls)129        # return [130        #     datasets.SplitGenerator(131        #         name=datasets.Split.TRAIN,132        #         # These kwargs will be passed to _generate_examples133        #         gen_kwargs={134        #             "filepath": os.path.join(data_dir, "Shellcode_IA32.tsv"),135        #             "split": "train",136        #         },137        #     ),138        #     datasets.SplitGenerator(139        #         name=datasets.Split.TEST,140        #         # These kwargs will be passed to _generate_examples141        #         gen_kwargs={142        #             "filepath": os.path.join(data_dir, "Shellcode_IA32.tsv"),143        #             "split": "test"144        #         },145        #     ),146        #     datasets.SplitGenerator(147        #         name=datasets.Split.VALIDATION,148        #         # These kwargs will be passed to _generate_examples149        #         gen_kwargs={150        #             "filepath": os.path.join(data_dir, "Shellcode_IA32.tsv"),151        #             "split": "dev",152        #         },153        #     ),154        # ]155 156        return [157            datasets.SplitGenerator(158                name=datasets.Split.TRAIN,159                # These kwargs will be passed to _generate_examples160                gen_kwargs={161                    "filepath": os.path.join(data_dir),162                    "split": "train",163                },164            ),165            datasets.SplitGenerator(166                name=datasets.Split.TEST,167                # These kwargs will be passed to _generate_examples168                gen_kwargs={169                    "filepath": os.path.join(data_dir),170                    "split": "test"171                },172            ),173            datasets.SplitGenerator(174                name=datasets.Split.VALIDATION,175                # These kwargs will be passed to _generate_examples176                gen_kwargs={177                    "filepath": os.path.join(data_dir),178                    "split": "dev",179                },180            ),181        ]182 183    def _generate_examples(184        self, filepath, split  # method parameters are unpacked from `gen_kwargs` as given in `_split_generators`185    ):186        """ Yields examples as (key, example) tuples. """187        # This method handles input defined in _split_generators to yield (key, example) tuples from the dataset.188        # The `key` is here for legacy reason (tfds) and is not important in itself.189        """This function returns the examples in the raw (text) form."""190 191        df = pd.read_csv(filepath, delimiter = '\t')192        train = df.sample(frac = 0.8, random_state = 0)193        test = df.drop(train.index)194        dev = test.sample(frac = 0.5, random_state = 0)195        test = test.drop(dev.index)196 197        if split == 'train':198            data = train199        elif split == 'dev':200            data = dev201        elif split == 'test':202            data = test203 204        for idx, row in data.iterrows():205            yield idx, {206                "snippet": row["SNIPPETS"],207                "intent": row["INTENTS"],208 209            }210            # with open(filepath, encoding="utf-8") as f:211            #     reader = csv.DictReader(f, delimiter="\t", quoting=csv.QUOTE_NONE)212            #     reader =213            #     for idx, row in enumerate(reader):214            #215            #         yield idx, {216            #             "snippet": row["SNIPPETS"],217            #             "intent": row["INTENTS"],218            #219            #         }220