From 3a74ca87047420f1b4007f203472db9c0ec5a13d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= Date: Mon, 22 Jan 2024 13:45:25 +0100 Subject: [PATCH] chg: Update typing to please lief --- pymisp/tools/create_misp_object.py | 13 ++++++++----- pymisp/tools/elfobject.py | 7 ++++--- pymisp/tools/fileobject.py | 3 +-- pymisp/tools/machoobject.py | 9 +++++---- pymisp/tools/peobject.py | 7 ++++--- 5 files changed, 22 insertions(+), 17 deletions(-) diff --git a/pymisp/tools/create_misp_object.py b/pymisp/tools/create_misp_object.py index 7864905..7c2a181 100644 --- a/pymisp/tools/create_misp_object.py +++ b/pymisp/tools/create_misp_object.py @@ -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 diff --git a/pymisp/tools/elfobject.py b/pymisp/tools/elfobject.py index 664bc83..80ad7c8 100644 --- a/pymisp/tools/elfobject.py +++ b/pymisp/tools/elfobject.py @@ -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)}') diff --git a/pymisp/tools/fileobject.py b/pymisp/tools/fileobject.py index 696eeed..9a8c8be 100644 --- a/pymisp/tools/fileobject.py +++ b/pymisp/tools/fileobject.py @@ -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]") diff --git a/pymisp/tools/machoobject.py b/pymisp/tools/machoobject.py index 4a71db4..58f2e70 100644 --- a/pymisp/tools/machoobject.py +++ b/pymisp/tools/machoobject.py @@ -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)}') diff --git a/pymisp/tools/peobject.py b/pymisp/tools/peobject.py index 151ed3b..92b3ff8 100644 --- a/pymisp/tools/peobject.py +++ b/pymisp/tools/peobject.py @@ -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)}')