fix: Allow object-relation names with uppercase characters defined in the templates

pull/1075/head
Raphaël Vinot 2023-10-04 11:46:49 +02:00
parent 309aa62210
commit 5b0b4c65e2
1 changed files with 19 additions and 14 deletions

View File

@ -2,21 +2,25 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
from pathlib import Path from pathlib import Path
from typing import List, Optional
import csv import csv
from pymisp import MISPObject from pymisp import MISPObject
class CSVLoader(): 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 = '"'): delimiter: str = ',', quotechar: str = '"'):
self.template_name = template_name self.template_name = template_name
self.delimiter = delimiter self.delimiter = delimiter
self.quotechar = quotechar self.quotechar = quotechar
self.csv_path = csv_path 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 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 self.has_fieldnames = True
else: else:
self.has_fieldnames = has_fieldnames self.has_fieldnames = has_fieldnames
@ -28,22 +32,23 @@ class CSVLoader():
with open(self.csv_path, newline='') as csvfile: with open(self.csv_path, newline='') as csvfile:
reader = csv.reader(csvfile, delimiter=self.delimiter, quotechar=self.quotechar) reader = csv.reader(csvfile, delimiter=self.delimiter, quotechar=self.quotechar)
if self.has_fieldnames: if self.has_fieldnames:
# The file has fieldnames, we either ignore it, or validate its validity # The file has fieldnames, we either ignore it, or use them as object-relation
fieldnames = [f.strip().lower() for f in reader.__next__()] fieldnames = [f.strip() for f in reader.__next__()]
if not self.fieldnames: if not self.fieldnames:
self.fieldnames = fieldnames self.fieldnames = fieldnames
if not self.fieldnames: if not self.fieldnames:
raise Exception('No fieldnames, impossible to create objects.') 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 # Check if the CSV file has a header, and if it matches with the object template
tmp_object = MISPObject(self.template_name) 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.') if not tmp_object._definition['attributes']:
allowed_fieldnames = list(tmp_object._definition['attributes'].keys()) raise Exception(f'Unable to find the object template ({self.template_name}), impossible to create objects.')
for fieldname in self.fieldnames: allowed_fieldnames = list(tmp_object._definition['attributes'].keys())
if fieldname not in allowed_fieldnames: for fieldname in self.fieldnames:
raise Exception(f'{fieldname} is not a valid object relation for {self.template_name}: {allowed_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: for row in reader:
tmp_object = MISPObject(self.template_name) tmp_object = MISPObject(self.template_name)