From e0fac90310f820f7a5b6e7a1d1a683ba22c1a437 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= Date: Tue, 2 Jul 2019 11:55:51 +0200 Subject: [PATCH] new: Allow to pass delimiter & quotechar to the CSV loader --- examples/load_csv.py | 5 ++++- pymisp/tools/csvloader.py | 7 +++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/examples/load_csv.py b/examples/load_csv.py index 892dfbb..bad5d1a 100755 --- a/examples/load_csv.py +++ b/examples/load_csv.py @@ -35,6 +35,8 @@ if __name__ == '__main__': parser.add_argument("-f", "--fieldnames", nargs='*', default=[], help="Fieldnames of the CSV, have to match the object-relation allowed in the template. If empty, the fieldnames of the CSV have to match the template.") parser.add_argument("-s", "--skip_fieldnames", action='store_true', help="Skip fieldnames in the CSV.") parser.add_argument("-d", "--dump", action='store_true', help="(Debug) Dump the object in the terminal.") + parser.add_argument("--delimiter", type=str, default=',', help="Delimiter between firlds in the CSV. Default: ','.") + parser.add_argument("--quotechar", type=str, default='"', help="Quote character of the fields in the CSV. Default: '\"'.") # Interact with MISP misp_group = parser.add_mutually_exclusive_group() @@ -48,7 +50,8 @@ if __name__ == '__main__': else: has_fieldnames = args.skip_fieldnames csv_loader = CSVLoader(template_name=args.object_name, csv_path=args.path, - fieldnames=args.fieldnames, has_fieldnames=has_fieldnames) + fieldnames=args.fieldnames, has_fieldnames=has_fieldnames, + delimiter=args.delimiter, quotechar=args.quotechar) objects = csv_loader.load() if args.dump: diff --git a/pymisp/tools/csvloader.py b/pymisp/tools/csvloader.py index 86efe39..4dc44b7 100644 --- a/pymisp/tools/csvloader.py +++ b/pymisp/tools/csvloader.py @@ -8,8 +8,11 @@ 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: list=[], 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] if not self.fieldnames: @@ -23,7 +26,7 @@ class CSVLoader(): objects = [] with open(self.csv_path, newline='') as csvfile: - reader = csv.reader(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__()]