chg: Use pathlib, cleanup

pull/7/head
Raphaël Vinot 2019-07-25 17:11:52 +02:00
parent edba867c35
commit 464a001aa4
1 changed files with 8 additions and 8 deletions

View File

@ -3,11 +3,11 @@
import json import json
from json import JSONEncoder from json import JSONEncoder
import os
import sys import sys
import collections import collections
from glob import glob from glob import glob
from ipaddress import ip_address, ip_network from ipaddress import ip_address, ip_network
from pathlib import Path
try: try:
@ -41,7 +41,7 @@ class WarningList():
self.version = int(self.warninglist['version']) self.version = int(self.warninglist['version'])
self.name = self.warninglist['name'] self.name = self.warninglist['name']
if self.warninglist['type'] not in self.expected_types: if self.warninglist['type'] not in self.expected_types:
raise Exception('Unexpected type, please update the expected_type list') raise PyMISPWarningListsError('Unexpected type, please update the expected_type list')
self.type = self.warninglist['type'] self.type = self.warninglist['type']
if self.warninglist.get('matching_attributes'): if self.warninglist.get('matching_attributes'):
self.matching_attributes = self.warninglist['matching_attributes'] self.matching_attributes = self.warninglist['matching_attributes']
@ -56,7 +56,7 @@ class WarningList():
self.slow_search = False self.slow_search = False
def __repr__(self): def __repr__(self):
return '<{self.__class__.__name__}(type="{self.name}", version="{self.version}", description="{self.description}")'.format(self=self) return f'<{self.__class__.__name__}(type="{self.name}", version="{self.version}", description="{self.description}")'
def __contains__(self, value): def __contains__(self, value):
if self.slow_search: if self.slow_search:
@ -118,11 +118,12 @@ class WarningLists(collections.Mapping):
""" """
if not lists: if not lists:
lists = [] lists = []
self.root_dir_warninglists = os.path.join(os.path.abspath(os.path.dirname(sys.modules['pymispwarninglists'].__file__)), self.root_dir_warninglists = Path(sys.modules['pymispwarninglists'].__file__).parent / 'data' / 'misp-warninglists' / 'lists'
'data', 'misp-warninglists', 'lists') for warninglist_file in glob(str(self.root_dir_warninglists / '*' / 'list.json')):
for warninglist_file in glob(os.path.join(self.root_dir_warninglists, '*', 'list.json')):
with open(warninglist_file, 'r') as f: with open(warninglist_file, 'r') as f:
lists.append(json.load(f)) lists.append(json.load(f))
if not lists:
raise PyMISPWarningListsError('Unable to load the lists. Do not forget to initialize the submodule (git submodule update --init).')
self.warninglists = {} self.warninglists = {}
for warninglist in lists: for warninglist in lists:
self.warninglists[warninglist['name']] = WarningList(warninglist, slow_search) self.warninglists[warninglist['name']] = WarningList(warninglist, slow_search)
@ -130,8 +131,7 @@ class WarningLists(collections.Mapping):
def validate_with_schema(self): def validate_with_schema(self):
if not HAS_JSONSCHEMA: if not HAS_JSONSCHEMA:
raise ImportError('jsonschema is required: pip install jsonschema') raise ImportError('jsonschema is required: pip install jsonschema')
schema = os.path.join(os.path.abspath(os.path.dirname(sys.modules['pymispwarninglists'].__file__)), schema = Path(sys.modules['pymispwarninglists'].__file__).parent / 'data' / 'misp-warninglists' / 'schema.json'
'data', 'misp-warninglists', 'schema.json')
with open(schema, 'r') as f: with open(schema, 'r') as f:
loaded_schema = json.load(f) loaded_schema = json.load(f)
for w in self.warninglists.values(): for w in self.warninglists.values():