Adds zlib-inflate, zlib-deflate, unmap-pe utils. Fixes bug in api_key…

… for elasticsearch settings

elastic/thrunting_tools/binaries/unmap_pe.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+"""CLI utility for unmapping PE memory regions"""
+# Licensed to Elasticsearch B.V. under one or more contributor
+# license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright
+# ownership. Elasticsearch B.V. licenses this file to you under
+# the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import logging
+from hashlib import sha256
+from pathlib import Path
+from typing import BinaryIO, Optional
+
+import pefile
+import typer
+
+from elastic.thrunting_tools.common.utils import stream, version_callback
+
+logger = logging.getLogger(__name__)
+MAX_CHUNK_SIZE: int = 4096
+
+app = typer.Typer(add_completion=False)
+
+
+@app.command()
+def unmap_pefile(
+    path_in: Path = typer.Option(
+        "-",
+        "--input",
+        "-i",
+        allow_dash=True,
+        readable=True,
+        file_okay=True,
+        dir_okay=False,
+        help="Filename for input stream",
+    ),
+    path_out: Path = typer.Option(
+        "-",
+        "--output",
+        "-o",
+        allow_dash=True,
+        writable=True,
+        file_okay=True,
+        dir_okay=False,
+        help="Filename for output stream",
+    ),
+    version: Optional[bool] = typer.Option(  # pylint: disable=unused-argument
+        None, "--version", callback=version_callback, help="Show version info and exit"
+    ),
+):
+    """
+    Process a PE file object, removing the memory mapping. This is useful when analyzing PE objects
+    captured from memory. Defaults to reading from standard in and writing to standard out.
+    """
+
+    f_in: BinaryIO
+    f_out: BinaryIO
+    shasum = sha256()
+
+    with stream(path_in, "rb") as f_in, stream(path_out, "wb") as f_out:
+        buff = bytearray()
+        while True:  # loop until EOF
+            _chunk: bytes = f_in.read(MAX_CHUNK_SIZE)
+            if not _chunk:  # An empty string is the end
+                break
+            buff.extend(_chunk)
+            shasum.update(_chunk)
+
+        pe: pefile.PE
+        try:
+            pe = pefile.PE(data=buff)
+        except pefile.PEFormatError:
+            logger.error("Unable to process file as PE. sha256: %s", shasum.hexdigest())
+            raise typer.Exit(1)  # pylint: disable=raise-missing-from
+
+        header_len: int = pe.OPTIONAL_HEADER.SizeOfHeaders
+        f_out.write(buff[:header_len])
+
+        entry: pefile.SectionStructure
+        for entry in pe.sections:
+            f_out.write(
+                buff[entry.VirtualAddress : entry.VirtualAddress + entry.SizeOfRawData]
+            )
+
+
+if __name__ == "__main__":
+    app()

elastic/thrunting_tools/common/elastic.py

@@ -40,7 +40,7 @@ def connect_elasticsearch(settings: ElasticsearchSettings) -> Elasticsearch:
         )
 
     if settings.api_key:
-        _apikey = tuple(settings.get("api_key").split(":"))
+        _apikey = tuple(settings.api_key.split(":"))
 
     if settings.cloud_id:
         logger.debug("Connecting to Elasticsearch using cloud_id %s", settings.cloud_id)

elastic/thrunting_tools/common/utils.py

@@ -16,10 +16,13 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import sys
+from collections.abc import Generator
+from contextlib import contextmanager
 from importlib.metadata import version
 from logging import getLogger
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, BinaryIO, Dict, List, TextIO
 
 import typer
 from ruamel.yaml import YAML
@@ -63,3 +66,18 @@ def version_callback(value: bool):
         print(f"Elastic Security Labs Thrunting Tools, {_version}")
         print("https://github.com/elastic/securitylabs-thrunting-tools")
         raise typer.Exit()
+
+
+@contextmanager
+def stream(arg: Path, mode: str = "r") -> Generator[TextIO | BinaryIO, None, None]:
+    """Provides a context manager that handles stdin/stdout as well"""
+    if mode not in ("r", "w", "rb", "wb"):
+        raise ValueError('mode not "r", "w", "rb", "wb"')
+    if str(arg) == "-":
+        if mode in ("r", "w"):
+            yield sys.stdin if mode == "r" else sys.stdout
+        else:
+            yield sys.stdin.buffer if mode == "rb" else sys.stdout.buffer
+    else:
+        with arg.open(mode) as handle:
+            yield handle

elastic/thrunting_tools/compression/zlib_deflate.py

@@ -6,17 +6,17 @@
 
 import typer
 
-from elastic.thrunting_tools.common.utils import version_callback
+from elastic.thrunting_tools.common.utils import stream, version_callback
 
 MAX_CHUNK_SIZE: int = 4096
 
 app = typer.Typer(add_completion=False)
 
 
-def chunk_generator(stream: BinaryIO) -> str:
+def chunk_generator(stream_in: BinaryIO) -> str:
     _deflator = zlib.compressobj()
     while True:  # Loop until EOF
-        _chunk = stream.read(MAX_CHUNK_SIZE)
+        _chunk = stream_in.read(MAX_CHUNK_SIZE)
         if not _chunk:
             yield _deflator.flush()
             break
@@ -50,12 +50,10 @@ def zlib_deflate(
     ),
 ):
 
-    f_in = sys.stdin.buffer if str(path_in) == "-" else path_in.open("rb")
-    f_out = sys.stdout.buffer if str(path_out) == "-" else path_out.open("wb")
-
-    gen = chunk_generator(f_in)
-    for chunk in gen:
-        f_out.write(chunk)
+    with stream(path_in, "rb") as f_in, stream(path_out, "wb") as f_out:
+        gen = chunk_generator(f_in)
+        for chunk in gen:
+            f_out.write(chunk)
 
 
 if __name__ == "__main__":

elastic/thrunting_tools/compression/zlib_inflate.py

@@ -1,22 +1,21 @@
 #!/usr/bin/env python3
-import sys
 import zlib
 from pathlib import Path
 from typing import BinaryIO, Optional
 
 import typer
 
-from elastic.thrunting_tools.common.utils import version_callback
+from elastic.thrunting_tools.common.utils import stream, version_callback
 
 MAX_CHUNK_SIZE: int = 4096
 
 app = typer.Typer(add_completion=False)
 
 
-def chunk_generator(stream: BinaryIO) -> str:
+def chunk_generator(stream_in: BinaryIO) -> str:
     _inflator = zlib.decompressobj()
     while True:  # Loop until EOF
-        _chunk = stream.read(MAX_CHUNK_SIZE)
+        _chunk = stream_in.read(MAX_CHUNK_SIZE)
         if not _chunk:  # an empty string is the end
             yield _inflator.flush()
             break
@@ -49,13 +48,14 @@ def zlib_inflate(
         None, "--version", callback=version_callback, help="Show version info and exit"
     ),
 ):
-
-    f_in = sys.stdin.buffer if str(path_in) == "-" else path_in.open("rb")
-    f_out = sys.stdout.buffer if str(path_out) == "-" else path_out.open("wb")
-
-    gen = chunk_generator(f_in)
-    for chunk in gen:
-        f_out.write(chunk)
+    """
+    Decompresses (inflates) a zlib compressed file. Defaults to reading from standard
+    in and writing to standard out.
+    """
+    with stream(path_in, "rb") as f_in, stream(path_out, "wb") as f_out:
+        gen = chunk_generator(f_in)
+        for chunk in gen:
+            f_out.write(chunk)
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -25,10 +25,13 @@ repository = "https://github.com/elastic/securitylabs-thrunting-tools"
 [tool.poetry.scripts]
 eql-query       = 'elastic.thrunting_tools.eql_query:app'
 lucene-query    = 'elastic.thrunting_tools.lucene_query:app'
+unmap-pe        = 'elastic.thrunting_tools.binaries.unmap_pe:app'
 from-charcode   = 'elastic.thrunting_tools.format.from_charcode:app'
 to-charcode     = 'elastic.thrunting_tools.format.to_charcode:app'
 url-decode      = 'elastic.thrunting_tools.format.url_decode:app'
 url-encode      = 'elastic.thrunting_tools.format.url_encode:app'
+zlib-inflate    = 'elastic.thrunting_tools.compression.zlib_inflate:app'
+zlib-deflate    = 'elastic.thrunting_tools.compression.zlib_deflate:app'
 zlib-compress   = 'elastic.thrunting_tools.compression.zlib_deflate:app'
 zlib-decompress = 'elastic.thrunting_tools.compression.zlib_inflate:app'
 
@@ -43,6 +46,7 @@ elasticsearch  = "^8.5.0"
 appdirs        = "^1.4.4"
 ruamel-yaml    = "^0.17.21"
 pydantic       = { version = "^1.10.2" }
+pefile         = "^2022.5.30"
 
 [tool.poetry.group.dev.dependencies]
 pytest     = "^7.2.0"