flow.record 3.19.dev3__tar.gz → 3.19.dev5__tar.gz

{flow_record-3.19.dev3/flow.record.egg-info → flow_record-3.19.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: flow.record
-Version: 3.19.dev3
+Version: 3.19.dev5
 Summary: A library for defining and creating structured data (called records) that can be streamed to disk or piped to other tools that use flow.record
 Author-email: Dissect Team <dissect@fox-it.com>
 License: Affero General Public License v3

{flow_record-3.19.dev3 → flow_record-3.19.dev5}/flow/record/__init__.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import gzip
-import os
+from pathlib import Path
 
 from flow.record.base import (
     IGNORE_FIELDS_FOR_COMPARISON,
@@ -39,71 +41,61 @@ from flow.record.stream import (
 
 __all__ = [
     "IGNORE_FIELDS_FOR_COMPARISON",
-    "RECORD_VERSION",
     "RECORDSTREAM_MAGIC",
+    "RECORD_VERSION",
+    "DynamicDescriptor",
     "FieldType",
-    "Record",
     "GroupedRecord",
-    "RecordDescriptor",
+    "JsonRecordPacker",
+    "PathTemplateWriter",
+    "Record",
     "RecordAdapter",
+    "RecordArchiver",
+    "RecordDescriptor",
+    "RecordDescriptorError",
     "RecordField",
-    "RecordReader",
-    "RecordWriter",
     "RecordOutput",
-    "RecordPrinter",
     "RecordPacker",
-    "JsonRecordPacker",
-    "RecordStreamWriter",
+    "RecordPrinter",
+    "RecordReader",
     "RecordStreamReader",
-    "open_path_or_stream",
+    "RecordStreamWriter",
+    "RecordWriter",
+    "dynamic_fieldtype",
+    "extend_record",
+    "ignore_fields_for_comparison",
+    "iter_timestamped_records",
     "open_path",
+    "open_path_or_stream",
     "open_stream",
-    "ignore_fields_for_comparison",
+    "record_stream",
     "set_ignored_fields_for_comparison",
     "stream",
-    "dynamic_fieldtype",
-    "DynamicDescriptor",
-    "PathTemplateWriter",
-    "RecordArchiver",
-    "RecordDescriptorError",
-    "record_stream",
-    "extend_record",
-    "iter_timestamped_records",
 ]
 
 
-class View:
-    fields = None
-
-    def __init__(self, fields):
-        self.fields = fields
-
-    def __iter__(self, fields):
-        pass
-
-
 class RecordDateSplitter:
     basepath = None
     out = None
 
-    def __init__(self, basepath):
-        self.basepath = basepath
+    def __init__(self, basepath: str | Path):
+        self.basepath = Path(basepath)
         self.out = {}
 
-    def getstream(self, t):
+    def getstream(self, t: tuple[int, int, int]) -> RecordStreamWriter:
         if t not in self.out:
-            path = os.path.join(self.basepath, "-".join(["{:2d}".format(v) for v in t]) + ".rec.gz")
+            path = self.basepath.joinpath("-".join([f"{v:2d}" for v in t]) + ".rec.gz")
             f = gzip.GzipFile(path, "wb")
             rs = RecordStreamWriter(f)
             self.out[t] = rs
         return self.out[t]
 
-    def write(self, r):
+    def write(self, r: Record) -> None:
         t = (r.ts.year, r.ts.month, r.ts.day)
         rs = self.getstream(t)
         rs.write(r)
         rs.fp.flush()
 
-    def close(self):
+    def close(self) -> None:
         for rs in self.out.values():
             rs.close()

flow_record-3.19.dev5/flow/record/adapter/__init__.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # make this namespace extensible from other packages
+import abc
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from flow.record.base import Record
+
+
+class AbstractWriter(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def write(self, rec: Record) -> None:
+        """Write a record."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def flush(self) -> None:
+        """Flush any buffered writes."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def close(self) -> None:
+        """Close the Writer, no more writes will be possible."""
+        raise NotImplementedError
+
+    def __del__(self) -> None:
+        self.close()
+
+    def __enter__(self) -> AbstractWriter:  # noqa: PYI034
+        return self
+
+    def __exit__(self, *args) -> None:
+        self.flush()
+        self.close()
+
+
+class AbstractReader(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def __iter__(self) -> Iterator[Record]:
+        """Return a record iterator."""
+        raise NotImplementedError
+
+    def close(self) -> None:  # noqa: B027
+        """Close the Reader, can be overriden to properly free resources."""
+
+    def __enter__(self) -> AbstractReader:  # noqa: PYI034
+        return self
+
+    def __exit__(self, *args) -> None:
+        self.close()

{flow_record-3.19.dev3 → flow_record-3.19.dev5}/flow/record/adapter/archive.py

@@ -1,6 +1,13 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from flow.record.adapter import AbstractReader, AbstractWriter
 from flow.record.stream import RecordArchiver
 
+if TYPE_CHECKING:
+    from flow.record.base import Record
+
 __usage__ = """
 Record archiver adapter, writes records to YYYY/mm/dd directories (writer only)
 ---
@@ -12,7 +19,7 @@ Write usage: rdump -w archive://[PATH]
 class ArchiveWriter(AbstractWriter):
     writer = None
 
-    def __init__(self, path, **kwargs):
+    def __init__(self, path: str, **kwargs):
         self.path = path
 
         path_template = kwargs.get("path_template")
@@ -20,19 +27,19 @@ class ArchiveWriter(AbstractWriter):
 
         self.writer = RecordArchiver(self.path, path_template=path_template, name=name)
 
-    def write(self, r):
+    def write(self, r: Record) -> None:
         self.writer.write(r)
 
-    def flush(self):
+    def flush(self) -> None:
         # RecordArchiver already flushes after every write
         pass
 
-    def close(self):
+    def close(self) -> None:
         if self.writer:
             self.writer.close()
         self.writer = None
 
 
 class ArchiveReader(AbstractReader):
-    def __init__(self, path, **kwargs):
+    def __init__(self, path: str, **kwargs):
         raise NotImplementedError

{flow_record-3.19.dev3 → flow_record-3.19.dev5}/flow/record/adapter/avro.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import json
 from datetime import datetime, timedelta, timezone
 from importlib.util import find_spec
-from typing import Any, Iterator
+from typing import TYPE_CHECKING, Any, BinaryIO
 
 import fastavro
 
@@ -12,6 +12,10 @@ from flow.record.adapter import AbstractReader, AbstractWriter
 from flow.record.selector import make_selector
 from flow.record.utils import is_stdout
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from pathlib import Path
+
 __usage__ = """
 Apache AVRO adapter
 ---
@@ -52,7 +56,7 @@ class AvroWriter(AbstractWriter):
     fp = None
     writer = None
 
-    def __init__(self, path, key=None, **kwargs):
+    def __init__(self, path: str | Path | BinaryIO, **kwargs):
         self.fp = record.open_path_or_stream(path, "wb")
 
         self.desc = None
@@ -69,11 +73,11 @@ class AvroWriter(AbstractWriter):
             self.writer = fastavro.write.Writer(self.fp, self.parsed_schema, codec=self.codec)
 
         if self.desc != r._desc:
-            raise Exception("Mixed record types")
+            raise ValueError("Mixed record types")
 
         self.writer.write(r._packdict())
 
-    def flush(self):
+    def flush(self) -> None:
         if not self.writer:
             self.writer = fastavro.write.Writer(
                 self.fp,
@@ -92,21 +96,21 @@ class AvroWriter(AbstractWriter):
 class AvroReader(AbstractReader):
     fp = None
 
-    def __init__(self, path, selector=None, **kwargs):
+    def __init__(self, path: str, selector: str | None = None, **kwargs):
         self.fp = record.open_path_or_stream(path, "rb")
         self.selector = make_selector(selector)
 
         self.reader = fastavro.reader(self.fp)
         self.schema = self.reader.writer_schema
         if not self.schema:
-            raise Exception("Missing Avro schema")
+            raise ValueError("Missing Avro schema")
 
         self.desc = schema_to_descriptor(self.schema)
 
         # Store the fieldnames that are of type "datetime"
-        self.datetime_fields = set(
+        self.datetime_fields = {
             name for name, field in self.desc.get_all_fields().items() if field.typename == "datetime"
-        )
+        }
 
     def __iter__(self) -> Iterator[record.Record]:
         for obj in self.reader:
@@ -149,7 +153,7 @@ def descriptor_to_schema(desc: record.RecordDescriptor) -> dict[str, Any]:
         else:
             avro_type = AVRO_TYPE_MAP.get(field_type)
             if not avro_type:
-                raise Exception("Unsupported Avro type: {}".format(field_type))
+                raise ValueError(f"Unsupported Avro type: {field_type}")
 
             field_schema["type"] = [avro_type, "null"]
 
@@ -190,11 +194,10 @@ def avro_type_to_flow_type(ftype: list) -> str:
         if isinstance(t, dict):
             if t.get("type") == "array":
                 item_type = avro_type_to_flow_type(t.get("items"))
-                return "{}[]".format(item_type)
-            else:
-                logical_type = t.get("logicalType")
-                if logical_type and ("time" in logical_type or "date" in logical_type):
-                    return "datetime"
+                return f"{item_type}[]"
+            logical_type = t.get("logicalType")
+            if logical_type and ("time" in logical_type or "date" in logical_type):
+                return "datetime"
 
         if t == "null":
             continue
@@ -202,4 +205,4 @@ def avro_type_to_flow_type(ftype: list) -> str:
         if t in RECORD_TYPE_MAP:
             return RECORD_TYPE_MAP[t]
 
-    raise TypeError("Can't map avro type to flow type: {}".format(t))
+    raise TypeError(f"Can't map avro type to flow type: {t}")

{flow_record-3.19.dev3 → flow_record-3.19.dev5}/flow/record/adapter/broker.py

@@ -1,7 +1,15 @@
-from flow.broker import Publisher, Subscriber
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
+from flow.broker import Publisher, Subscriber
 from flow.record.adapter import AbstractReader, AbstractWriter
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from flow.record.base import Record
+
 __usage__ = """
 PubSub adapter using flow.broker
 ---
@@ -13,23 +21,23 @@ Read usage: rdump broker+tcp://[IP]:[PORT] -s True
 class BrokerWriter(AbstractWriter):
     publisher = None
 
-    def __init__(self, uri, source=None, classification=None, **kwargs):
+    def __init__(self, uri: str, source: str | None = None, classification: str | None = None, **kwargs):
         self.publisher = Publisher(uri, **kwargs)
         self.source = source
         self.classification = classification
 
-    def write(self, r):
+    def write(self, r: Record) -> None:
         record = r._replace(
             _source=self.source or r._source,
             _classification=self.classification or r._classification,
         )
         self.publisher.send(record)
 
-    def flush(self):
+    def flush(self) -> None:
         if self.publisher:
             self.publisher.flush()
 
-    def close(self):
+    def close(self) -> None:
         if self.publisher:
             if hasattr(self.publisher, "stop"):
                 # Requires flow.broker >= 1.1.1
@@ -42,14 +50,14 @@ class BrokerWriter(AbstractWriter):
 class BrokerReader(AbstractReader):
     subscriber = None
 
-    def __init__(self, uri, name=None, selector=None, **kwargs):
+    def __init__(self, uri: str, name: str | None = None, selector: str | None = None, **kwargs):
         self.subscriber = Subscriber(uri, **kwargs)
         self.subscription = self.subscriber.select(name, str(selector))
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Record]:
         return iter(self.subscription)
 
-    def close(self):
+    def close(self) -> None:
         if self.subscriber:
             self.subscriber.stop()
         self.subscriber = None

{flow_record-3.19.dev3 → flow_record-3.19.dev5}/flow/record/adapter/csvfile.py

@@ -1,14 +1,19 @@
-from __future__ import absolute_import
+from __future__ import annotations
 
 import csv
 import sys
+from pathlib import Path
+from typing import TYPE_CHECKING
 
 from flow.record import RecordDescriptor
 from flow.record.adapter import AbstractReader, AbstractWriter
-from flow.record.base import normalize_fieldname
+from flow.record.base import Record, normalize_fieldname
 from flow.record.selector import make_selector
 from flow.record.utils import is_stdout
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 __usage__ = """
 Comma-separated values (CSV) adapter
 ---
@@ -23,13 +28,20 @@ Optional parameters:
 
 
 class CsvfileWriter(AbstractWriter):
-    def __init__(self, path, fields=None, exclude=None, lineterminator=None, **kwargs):
+    def __init__(
+        self,
+        path: str | Path | None,
+        fields: str | list[str] | None = None,
+        exclude: str | list[str] | None = None,
+        lineterminator: str = "\r\n",
+        **kwargs,
+    ):
         self.fp = None
         if path in (None, "", "-"):
             self.fp = sys.stdout
         else:
-            self.fp = open(path, "w", newline="")
-        self.lineterminator = lineterminator or "\r\n"
+            self.fp = Path(path).open("w", newline="")  # noqa: SIM115
+        self.lineterminator = lineterminator
         for r, n in ((r"\r", "\r"), (r"\n", "\n"), (r"\t", "\t")):
             self.lineterminator = self.lineterminator.replace(r, n)
         self.desc = None
@@ -41,7 +53,7 @@ class CsvfileWriter(AbstractWriter):
         if isinstance(self.exclude, str):
             self.exclude = self.exclude.split(",")
 
-    def write(self, r):
+    def write(self, r: Record) -> None:
         rdict = r._asdict(fields=self.fields, exclude=self.exclude)
         if not self.desc or self.desc != r._desc:
             self.desc = r._desc
@@ -49,24 +61,26 @@ class CsvfileWriter(AbstractWriter):
             self.writer.writeheader()
         self.writer.writerow(rdict)
 
-    def flush(self):
+    def flush(self) -> None:
         if self.fp:
             self.fp.flush()
 
-    def close(self):
+    def close(self) -> None:
         if self.fp and not is_stdout(self.fp):
             self.fp.close()
         self.fp = None
 
 
 class CsvfileReader(AbstractReader):
-    def __init__(self, path, selector=None, fields=None, **kwargs):
+    def __init__(
+        self, path: str | Path | None, selector: str | None = None, fields: str | list[str] | None = None, **kwargs
+    ):
         self.fp = None
         self.selector = make_selector(selector)
         if path in (None, "", "-"):
             self.fp = sys.stdin
         else:
-            self.fp = open(path, "r", newline="")
+            self.fp = Path(path).open("r", newline="")  # noqa: SIM115
 
         self.dialect = "excel"
         if self.fp.seekable():
@@ -87,12 +101,12 @@ class CsvfileReader(AbstractReader):
         # Create RecordDescriptor from fields, skipping fields starting with "_" (reserved for internal use)
         self.desc = RecordDescriptor("csv/reader", [("string", col) for col in self.fields if not col.startswith("_")])
 
-    def close(self):
+    def close(self) -> None:
         if self.fp:
             self.fp.close()
         self.fp = None
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Record]:
         for row in self.reader:
             rdict = dict(zip(self.fields, row))
             record = self.desc.init_from_dict(rdict)

{flow_record-3.19.dev3 → flow_record-3.19.dev5}/flow/record/adapter/elastic.py

@@ -4,7 +4,7 @@ import hashlib
 import logging
 import queue
 import threading
-from typing import Iterator
+from typing import TYPE_CHECKING
 
 import elasticsearch
 import elasticsearch.helpers
@@ -13,7 +13,11 @@ from flow.record.adapter import AbstractReader, AbstractWriter
 from flow.record.base import Record, RecordDescriptor
 from flow.record.fieldtypes import fieldtype_for_value
 from flow.record.jsonpacker import JsonRecordPacker
-from flow.record.selector import CompiledSelector, Selector
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from flow.record.selector import CompiledSelector, Selector
 
 __usage__ = """
 ElasticSearch adapter
@@ -25,6 +29,7 @@ Read usage: rdump elastic+[PROTOCOL]://[IP]:[PORT]?index=[INDEX]
 
 Optional arguments:
   [API_KEY]: base64 encoded api key to authenticate with (default: False)
+  [QUEUE_SIZE]: maximum queue size for writing records; limits memory usage (default: 100000)
   [INDEX]: name of the index to use (default: records)
   [VERIFY_CERTS]: verify certs of Elasticsearch instance (default: True)
   [HASH_RECORD]: make record unique by hashing record [slow] (default: False)
@@ -43,6 +48,7 @@ class ElasticWriter(AbstractWriter):
         http_compress: str | bool = True,
         hash_record: str | bool = False,
         api_key: str | None = None,
+        queue_size: int = 100000,
         **kwargs,
     ) -> None:
         self.index = index
@@ -50,11 +56,12 @@ class ElasticWriter(AbstractWriter):
         verify_certs = str(verify_certs).lower() in ("1", "true")
         http_compress = str(http_compress).lower() in ("1", "true")
         self.hash_record = str(hash_record).lower() in ("1", "true")
+        queue_size = int(queue_size)
 
         if not uri.lower().startswith(("http://", "https://")):
             uri = "http://" + uri
 
-        self.queue: queue.Queue[Record | StopIteration] = queue.Queue()
+        self.queue: queue.Queue[Record | StopIteration] = queue.Queue(maxsize=queue_size)
         self.event = threading.Event()
 
         self.es = elasticsearch.Elasticsearch(
@@ -147,7 +154,7 @@ class ElasticWriter(AbstractWriter):
         self.event.set()
 
     def write(self, record: Record) -> None:
-        self.queue.put_nowait(record)
+        self.queue.put(record)
 
     def flush(self) -> None:
         pass

{flow_record-3.19.dev3 → flow_record-3.19.dev5}/flow/record/adapter/jsonfile.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import json
+from typing import TYPE_CHECKING, BinaryIO
 
 from flow import record
 from flow.record import JsonRecordPacker
@@ -7,6 +10,12 @@ from flow.record.fieldtypes import fieldtype_for_value
 from flow.record.selector import make_selector
 from flow.record.utils import is_stdout
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from pathlib import Path
+
+    from flow.record.base import Record, RecordDescriptor
+
 __usage__ = """
 JSON adapter
 ---
@@ -21,7 +30,9 @@ Read usage: rdump jsonfile://[PATH]
 class JsonfileWriter(AbstractWriter):
     fp = None
 
-    def __init__(self, path, indent=None, descriptors=True, **kwargs):
+    def __init__(
+        self, path: str | Path | BinaryIO, indent: str | int | None = None, descriptors: bool = True, **kwargs
+    ):
         self.descriptors = str(descriptors).lower() in ("true", "1")
         self.fp = record.open_path_or_stream(path, "w")
         if isinstance(indent, str):
@@ -30,21 +41,21 @@ class JsonfileWriter(AbstractWriter):
         if self.descriptors:
             self.packer.on_descriptor.add_handler(self.packer_on_new_descriptor)
 
-    def packer_on_new_descriptor(self, descriptor):
+    def packer_on_new_descriptor(self, descriptor: RecordDescriptor) -> None:
         self._write(descriptor)
 
-    def _write(self, obj):
+    def _write(self, obj: Record | RecordDescriptor) -> None:
         record_json = self.packer.pack(obj)
         self.fp.write(record_json + "\n")
 
-    def write(self, r):
+    def write(self, r: Record) -> None:
         self._write(r)
 
-    def flush(self):
+    def flush(self) -> None:
         if self.fp:
             self.fp.flush()
 
-    def close(self):
+    def close(self) -> None:
         if self.fp and not is_stdout(self.fp):
             self.fp.close()
         self.fp = None
@@ -53,17 +64,17 @@ class JsonfileWriter(AbstractWriter):
 class JsonfileReader(AbstractReader):
     fp = None
 
-    def __init__(self, path, selector=None, **kwargs):
+    def __init__(self, path: str | Path | BinaryIO, selector: str | None = None, **kwargs):
         self.selector = make_selector(selector)
         self.fp = record.open_path_or_stream(path, "r")
         self.packer = JsonRecordPacker()
 
-    def close(self):
+    def close(self) -> None:
         if self.fp:
             self.fp.close()
         self.fp = None
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Record]:
         for line in self.fp:
             obj = self.packer.unpack(line)
             if isinstance(obj, record.Record):

{flow_record-3.19.dev3 → flow_record-3.19.dev5}/flow/record/adapter/line.py

@@ -60,12 +60,9 @@ class LineWriter(AbstractWriter):
         self.count += 1
         self.fp.write(f"--[ RECORD {self.count} ]--\n".encode())
         if rdict:
-            if rdict_types:
-                # also account for extra characters for fieldtype and whitespace + parenthesis
-                width = max(len(k + rdict_types[k]) for k in rdict) + 3
-            else:
-                width = max(len(k) for k in rdict)
-            fmt = "{{:>{width}}} = {{}}\n".format(width=width)
+            # also account for extra characters for fieldtype and whitespace + parenthesis
+            width = max(len(k + rdict_types[k]) for k in rdict) + 3 if rdict_types else max(len(k) for k in rdict)
+            fmt = f"{{:>{width}}} = {{}}\n"
         for key, value in rdict.items():
             if rdict_types:
                 key = f"{key} ({rdict_types[key]})"