From 1f589002d87a5c09729f253165489ba662eeca62 Mon Sep 17 00:00:00 2001 From: Stanislav Golovanov Date: Fri, 5 Jun 2026 22:50:18 +0300 Subject: [PATCH] Fix UnicodeEncodeError in csv adapter These can happen if record fields contain unicode surrogates when exporting to csv files. --- flow/record/adapter/csvfile.py | 7 ++++-- flow/record/fieldtypes/__init__.py | 16 +++++-------- flow/record/utils.py | 13 +++++++++++ tests/record/test_adapter.py | 37 ++++++++++++++++++++++++++++++ 4 files changed, 61 insertions(+), 12 deletions(-) diff --git a/flow/record/adapter/csvfile.py b/flow/record/adapter/csvfile.py index e920bc7e..3bab79bc 100644 --- a/flow/record/adapter/csvfile.py +++ b/flow/record/adapter/csvfile.py @@ -6,12 +6,12 @@ from pathlib import Path from typing import TYPE_CHECKING -from flow.record import RecordDescriptor +from flow.record import RecordDescriptor, fieldtypes from flow.record.adapter import AbstractReader, AbstractWriter from flow.record.base import normalize_fieldname from flow.record.context import get_app_context, match_record_with_context from flow.record.selector import make_selector -from flow.record.utils import boolean_argument, is_stdout +from flow.record.utils import boolean_argument, escape_surrogates, is_stdout if TYPE_CHECKING: from collections.abc import Iterator @@ -68,6 +68,9 @@ def write(self, r: Record) -> None: if self.header: # Write header only if it is requested self.writer.writeheader() + for k, v in rdict.items(): + if isinstance(v, fieldtypes.string): + rdict[k] = escape_surrogates(v) self.writer.writerow(rdict) def flush(self) -> None: diff --git a/flow/record/fieldtypes/__init__.py b/flow/record/fieldtypes/__init__.py index 4cd9d91c..1bfb6579 100644 --- a/flow/record/fieldtypes/__init__.py +++ b/flow/record/fieldtypes/__init__.py @@ -15,6 +15,8 @@ from typing import TYPE_CHECKING, Any from urllib.parse import urlparse +from flow.record.utils import escape_surrogates + try: try: from zoneinfo import ZoneInfo, ZoneInfoNotFoundError @@ -735,17 +737,11 @@ class posix_path(pathlib.PurePosixPath, path): class windows_path(pathlib.PureWindowsPath, path): + def __str__(self) -> str: + return escape_surrogates(super().__str__()) + def __repr__(self) -> str: - s = str(self) - # Only use repr() if we have surrogates that need escaping - try: - s.encode("utf-8") - except UnicodeEncodeError: - # Has surrogates - use repr but fix the over-escaping - s = repr(s)[1:-1] # This escapes surrogates as \udcXX - s = s.replace("\\\\", "\\") # Fix double backslashes - s = s.replace("\\'", "'") # Fix over-escaped quotes - s = s.replace('\\"', '"') # Fix over-escaped double quotes + s = escape_surrogates(str(self)) quote = "'" if "'" in s: diff --git a/flow/record/utils.py b/flow/record/utils.py index 89b3a3f2..1fe21a1b 100644 --- a/flow/record/utils.py +++ b/flow/record/utils.py @@ -86,6 +86,19 @@ def to_base64(value: str) -> str: return base64.b64encode(value).decode() +def escape_surrogates(s: str) -> str: + """Escape surrogate unicode symbols.""" + try: + s.encode("utf-8") + except UnicodeEncodeError: + # Has surrogates - use repr but fix the over-escaping + s = repr(s)[1:-1] # This escapes surrogates as \udcXX + s = s.replace("\\\\", "\\") # Fix double backslashes + s = s.replace("\\'", "'") # Fix over-escaped quotes + s = s.replace('\\"', '"') # Fix over-escaped double quotes + return s + + def catch_sigpipe(func: Callable[..., int]) -> Callable[..., int]: """Catches KeyboardInterrupt and BrokenPipeError (OSError 22 on Windows).""" diff --git a/tests/record/test_adapter.py b/tests/record/test_adapter.py index 6d49d311..ab7e3fca 100644 --- a/tests/record/test_adapter.py +++ b/tests/record/test_adapter.py @@ -27,6 +27,7 @@ LZ4_MAGIC, ZSTD_MAGIC, ) +from flow.record.fieldtypes import windows_path from flow.record.selector import CompiledSelector, Selector from tests._utils import generate_records @@ -455,6 +456,42 @@ def test_csv_adapter_lineterminator(capsysbinary: pytest.CaptureFixture) -> None assert out == b"count,foo,bar@0,hello,world@1,hello,world@2,hello,world@" +def test_csv_adapter_surrogates(capsysbinary: pytest.CaptureFixture) -> None: + TestRecord = RecordDescriptor( + "test/record", + [ + ("uint32", "count"), + ("string", "foo"), + ("string", "bar"), + ], + ) + + with RecordWriter(r"csvfile://?exclude=_source,_classification,_generated,_version") as writer: + rec = TestRecord(count=0, foo="hello", bar="world\udcce\udcc1\udcd9\udcc8") + writer.write(rec) + out, _ = capsysbinary.readouterr() + assert out == b"count,foo,bar\r\n0,hello,world\\udcce\\udcc1\\udcd9\\udcc8\r\n" + + +def test_csv_adapter_windows_path_surrogates(capsysbinary: pytest.CaptureFixture) -> None: + Record = RecordDescriptor( + "test/record", + [ + ("string", "name"), + ("path", "value"), + ], + ) + record = Record( + b"R\xc3\xa9\xeamy", + windows_path(b"\x43\x3a\x5c\xc3\xa4\xc3\x84\xe4".decode(errors="surrogateescape")), + ) + + with RecordWriter(r"csvfile://?exclude=_source,_classification,_generated,_version") as writer: + writer.write(record) + out, _ = capsysbinary.readouterr() + assert out.decode("utf-8") == "name,value\r\nRé\\udceamy,C:\\äÄ\\udce4\r\n" + + def test_csvfilereader(tmp_path: Path) -> None: path = tmp_path / "test.csv" with path.open("wb") as f: