Fix UnicodeEncodeError in csv adapter

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:

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:

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)."""
 

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: