From 5b0b4c65e2850ee515d197dbba464e21eeefb90c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= Date: Wed, 4 Oct 2023 11:46:49 +0200 Subject: [PATCH] fix: Allow object-relation names with uppercase characters defined in the templates --- pymisp/tools/csvloader.py | 33 +++++++++++++++++++-------------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/pymisp/tools/csvloader.py b/pymisp/tools/csvloader.py index b02ad69..c5880ac 100644 --- a/pymisp/tools/csvloader.py +++ b/pymisp/tools/csvloader.py @@ -2,21 +2,25 @@ # -*- coding: utf-8 -*- from pathlib import Path +from typing import List, Optional + import csv from pymisp import MISPObject class CSVLoader(): - def __init__(self, template_name: str, csv_path: Path, fieldnames: list = [], has_fieldnames=False, + def __init__(self, template_name: str, csv_path: Path, fieldnames: Optional[List[str]] = None, has_fieldnames=False, delimiter: str = ',', quotechar: str = '"'): self.template_name = template_name self.delimiter = delimiter self.quotechar = quotechar self.csv_path = csv_path - self.fieldnames = [f.strip().lower() for f in fieldnames] + self.fieldnames = [] + if fieldnames: + self.fieldnames = [f.strip() for f in fieldnames] if not self.fieldnames: - # If the user doesn't pass fieldnames, we assume the CSV has them. + # If the user doesn't pass fieldnames, they must be in the CSV. self.has_fieldnames = True else: self.has_fieldnames = has_fieldnames @@ -28,22 +32,23 @@ class CSVLoader(): with open(self.csv_path, newline='') as csvfile: reader = csv.reader(csvfile, delimiter=self.delimiter, quotechar=self.quotechar) if self.has_fieldnames: - # The file has fieldnames, we either ignore it, or validate its validity - fieldnames = [f.strip().lower() for f in reader.__next__()] + # The file has fieldnames, we either ignore it, or use them as object-relation + fieldnames = [f.strip() for f in reader.__next__()] if not self.fieldnames: self.fieldnames = fieldnames if not self.fieldnames: raise Exception('No fieldnames, impossible to create objects.') - else: - # Check if the CSV file has a header, and if it matches with the object template - tmp_object = MISPObject(self.template_name) - if not tmp_object._definition['attributes']: - raise Exception(f'Unable to find the object template ({self.template_name}), impossible to create objects.') - allowed_fieldnames = list(tmp_object._definition['attributes'].keys()) - for fieldname in self.fieldnames: - if fieldname not in allowed_fieldnames: - raise Exception(f'{fieldname} is not a valid object relation for {self.template_name}: {allowed_fieldnames}') + + # Check if the CSV file has a header, and if it matches with the object template + tmp_object = MISPObject(self.template_name) + + if not tmp_object._definition['attributes']: + raise Exception(f'Unable to find the object template ({self.template_name}), impossible to create objects.') + allowed_fieldnames = list(tmp_object._definition['attributes'].keys()) + for fieldname in self.fieldnames: + if fieldname not in allowed_fieldnames: + raise Exception(f'{fieldname} is not a valid object relation for {self.template_name}: {allowed_fieldnames}') for row in reader: tmp_object = MISPObject(self.template_name)