MISP
/
PyMISP
mirror of https://github.com/MISP/PyMISP
2
0
Fork 0
Code Issues Releases Wiki Activity

chg: Update typing to please lief

pull/1138/head
Raphaël Vinot 2024-01-22 13:45:25 +01:00
parent dbe29f87f3
commit 3a74ca8704
5 changed files with 22 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
pymisp/tools/create_misp_object.py
View File

@ -7,12 +7,12 @@ from io import BytesIO
from . import FileObject
from ..exceptions import MISPObjectException
import logging
from typing import Optional
logger = logging.getLogger('pymisp')
try:
import lief
import lief.logging
lief.logging.disable()
HAS_LIEF = True
@ -32,14 +32,17 @@ class FileTypeNotImplemented(MISPObjectException):
pass
def make_binary_objects(filepath: str | None = None, pseudofile: BytesIO | None = None, filename: str | None = None, standalone: bool = True, default_attributes_parameters: dict = {}):
def make_binary_objects(filepath: str | None = None, pseudofile: BytesIO | bytes | None = None, filename: str | None = None, standalone: bool = True, default_attributes_parameters: dict = {}):
misp_file = FileObject(filepath=filepath, pseudofile=pseudofile, filename=filename,
standalone=standalone, default_attributes_parameters=default_attributes_parameters)
if HAS_LIEF and (filepath or (pseudofile and filename)):
if HAS_LIEF and (filepath or pseudofile):
if filepath:
lief_parsed = lief.parse(filepath=filepath)
elif pseudofile and filename:
lief_parsed = lief.parse(raw=pseudofile.getvalue(), name=filename)
elif pseudofile:
if isinstance(pseudofile, bytes):
lief_parsed = lief.parse(raw=pseudofile)
else: # BytesIO
lief_parsed = lief.parse(obj=pseudofile)
else:
logger.critical('You need either a filepath, or a pseudofile and a filename.')
lief_parsed = None

7
pymisp/tools/elfobject.py
View File

@ -7,7 +7,6 @@ from ..exceptions import InvalidMISPObject
from io import BytesIO
from hashlib import md5, sha1, sha256, sha512
import logging
from typing import Union, Optional
from pathlib import Path
from . import FileObject
@ -33,15 +32,17 @@ def make_elf_objects(lief_parsed: lief.ELF.Binary, misp_file: FileObject, standa
class ELFObject(AbstractMISPObjectGenerator):
def __init__(self, parsed: lief.ELF.Binary | None = None, filepath: Path | str | None = None, pseudofile: BytesIO | None = None, **kwargs):
def __init__(self, parsed: lief.ELF.Binary | None = None, filepath: Path | str | None = None, pseudofile: BytesIO | bytes | list[int] | None = None, **kwargs):
"""Creates an ELF object, with lief"""
super().__init__('elf', **kwargs)
if not HAS_PYDEEP:
logger.warning("pydeep is missing, please install pymisp this way: pip install pymisp[fileobjects]")
if pseudofile:
if isinstance(pseudofile, BytesIO):
self.__elf = lief.ELF.parse(io=pseudofile)
self.__elf = lief.ELF.parse(obj=pseudofile)
elif isinstance(pseudofile, bytes):
self.__elf = lief.ELF.parse(raw=list(pseudofile))
elif isinstance(pseudofile, list):
self.__elf = lief.ELF.parse(raw=pseudofile)
else:
raise InvalidMISPObject(f'Pseudo file can be BytesIO or bytes got {type(pseudofile)}')

3
pymisp/tools/fileobject.py
View File

@ -11,7 +11,6 @@ import math
from collections import Counter
import logging
from pathlib import Path
from typing import Union, Optional
logger = logging.getLogger('pymisp')
@ -31,7 +30,7 @@ except ImportError:
class FileObject(AbstractMISPObjectGenerator):
def __init__(self, filepath: Path | str | None = None, pseudofile: BytesIO | None = None, filename: str | None = None, **kwargs) -> None:
def __init__(self, filepath: Path | str | None = None, pseudofile: BytesIO | bytes | None = None, filename: str | None = None, **kwargs) -> None:
super().__init__('file', **kwargs)
if not HAS_PYDEEP:
logger.warning("pydeep is missing, please install pymisp this way: pip install pymisp[fileobjects]")

9
pymisp/tools/machoobject.py
View File

@ -7,7 +7,6 @@ from .abstractgenerator import AbstractMISPObjectGenerator
from io import BytesIO
from hashlib import md5, sha1, sha256, sha512
import logging
from typing import Optional, Union
from pathlib import Path
from . import FileObject
@ -33,15 +32,17 @@ def make_macho_objects(lief_parsed: lief.MachO.Binary, misp_file: FileObject, st
class MachOObject(AbstractMISPObjectGenerator):
def __init__(self, parsed: lief.MachO.Binary | None = None, filepath: Path | str | None = None, pseudofile: BytesIO | None = None, **kwargs):
def __init__(self, parsed: lief.MachO.Binary | None = None, filepath: Path | str | None = None, pseudofile: BytesIO | list[int] | None = None, **kwargs):
"""Creates an MachO object, with lief"""
super().__init__('macho', **kwargs)
if not HAS_PYDEEP:
logger.warning("pydeep is missing, please install pymisp this way: pip install pymisp[fileobjects]")
if pseudofile:
if isinstance(pseudofile, BytesIO):
self.__macho = lief.MachO.parse(io=pseudofile)
self.__macho = lief.MachO.parse(obj=pseudofile)
elif isinstance(pseudofile, bytes):
self.__macho = lief.MachO.parse(raw=list(pseudofile))
elif isinstance(pseudofile, list):
self.__macho = lief.MachO.parse(raw=pseudofile)
else:
raise InvalidMISPObject(f'Pseudo file can be BytesIO or bytes got {type(pseudofile)}')
@ -49,7 +50,7 @@ class MachOObject(AbstractMISPObjectGenerator):
self.__macho = lief.MachO.parse(filepath)
elif parsed:
# Got an already parsed blob
if isinstance(parsed, lief.MachO.Binary):
if isinstance(parsed, lief.MachO.FatBinary):
self.__macho = parsed
else:
raise InvalidMISPObject(f'Not a lief.MachO.Binary: {type(parsed)}')

7
pymisp/tools/peobject.py
View File

@ -8,7 +8,6 @@ from io import BytesIO
from hashlib import md5, sha1, sha256, sha512
from datetime import datetime
import logging
from typing import Optional, Union
from pathlib import Path
from base64 import b64encode
@ -36,15 +35,17 @@ def make_pe_objects(lief_parsed: lief.PE.Binary, misp_file: FileObject, standalo
class PEObject(AbstractMISPObjectGenerator):
def __init__(self, parsed: lief.PE.Binary | None = None, filepath: Path | str | None = None, pseudofile: BytesIO | None = None, **kwargs):
def __init__(self, parsed: lief.PE.Binary | None = None, filepath: Path | str | None = None, pseudofile: BytesIO | list[int] | None = None, **kwargs):
"""Creates an PE object, with lief"""
super().__init__('pe', **kwargs)
if not HAS_PYDEEP:
logger.warning("pydeep is missing, please install pymisp this way: pip install pymisp[fileobjects]")
if pseudofile:
if isinstance(pseudofile, BytesIO):
self.__pe = lief.PE.parse(io=pseudofile)
self.__pe = lief.PE.parse(obj=pseudofile)
elif isinstance(pseudofile, bytes):
self.__pe = lief.PE.parse(raw=list(pseudofile))
elif isinstance(pseudofile, list):
self.__pe = lief.PE.parse(raw=pseudofile)
else:
raise InvalidMISPObject(f'Pseudo file can be BytesIO or bytes got {type(pseudofile)}')