diff --git a/packages/syft-restrict/src/syft_restrict/__init__.py b/packages/syft-restrict/src/syft_restrict/__init__.py index 947af1732b3..a565088f714 100644 --- a/packages/syft-restrict/src/syft_restrict/__init__.py +++ b/packages/syft-restrict/src/syft_restrict/__init__.py @@ -10,7 +10,13 @@ __version__ = "0.1.0" from .auditor import AuditReport, PathAudit, audit_allow_functions -from .errors import MarkerError, PolicyViolation, RestrictError, SourceEncodingError +from .errors import ( + MarkerError, + ObfuscationError, + PolicyViolation, + RestrictError, + SourceEncodingError, +) from .markers import parse_markers from .obfuscator import obfuscate from .policy import Policy @@ -32,6 +38,7 @@ "PolicyViolation", "MarkerError", "SourceEncodingError", + "ObfuscationError", "RestrictError", "__version__", ] diff --git a/packages/syft-restrict/src/syft_restrict/astutil.py b/packages/syft-restrict/src/syft_restrict/astutil.py index 3e0ca3e790b..75269d20756 100644 --- a/packages/syft-restrict/src/syft_restrict/astutil.py +++ b/packages/syft-restrict/src/syft_restrict/astutil.py @@ -3,6 +3,7 @@ from __future__ import annotations import ast +import io from pydantic import BaseModel @@ -22,6 +23,15 @@ def normalize_ranges(private) -> list[tuple[int, int]]: return ranges +def source_rows(text: str) -> list[str]: + """Split ``text`` into rows on ``\\n`` only, as ``tokenize`` counts them. + + Each row keeps its ``\\n``. Unlike ``str.splitlines()``, a form feed or another separator + character does not end a row. + """ + return io.StringIO(text).readlines() + + def row_in_ranges(row: int, ranges) -> bool: """True if a 1-based line number falls inside any range.""" return any(lo <= row <= hi for lo, hi in ranges) diff --git a/packages/syft-restrict/src/syft_restrict/errors.py b/packages/syft-restrict/src/syft_restrict/errors.py index b1fda724a78..ea5c1c0a523 100644 --- a/packages/syft-restrict/src/syft_restrict/errors.py +++ b/packages/syft-restrict/src/syft_restrict/errors.py @@ -30,3 +30,7 @@ class MarkerError(RestrictError): class SourceEncodingError(RestrictError): """Raised when a source file is not plain UTF-8.""" + + +class ObfuscationError(RestrictError): + """Raised when the display copy does not keep the rows of the source.""" diff --git a/packages/syft-restrict/src/syft_restrict/markers.py b/packages/syft-restrict/src/syft_restrict/markers.py index 081fbe392ef..4b2d9e61f4b 100644 --- a/packages/syft-restrict/src/syft_restrict/markers.py +++ b/packages/syft-restrict/src/syft_restrict/markers.py @@ -24,6 +24,7 @@ def attention(x): import re import tokenize +from .astutil import source_rows from .errors import MarkerError __all__ = ["parse_markers"] @@ -116,7 +117,7 @@ def __init__(self) -> None: def feed(self, source: str) -> None: markers = _scan_markers(source) - for line in range(1, len(source.splitlines()) + 1): + for line in range(1, len(source_rows(source)) + 1): marker = markers.get(line) if marker is None: if self.stack: diff --git a/packages/syft-restrict/src/syft_restrict/obfuscator.py b/packages/syft-restrict/src/syft_restrict/obfuscator.py index f9bfe838bea..732975d7cfe 100644 --- a/packages/syft-restrict/src/syft_restrict/obfuscator.py +++ b/packages/syft-restrict/src/syft_restrict/obfuscator.py @@ -23,6 +23,7 @@ node_in_ranges, normalize_ranges, row_in_ranges, + source_rows, ) from .policy import DEFAULT_KEEP @@ -112,13 +113,16 @@ def obfuscate(source: str, obfuscate_ranges, hide_ranges, scan: FileScan) -> str if new is not None: edits.append((srow, scol, erow, ecol, new)) elif tok.type == tokenize.STRING: - edits.append((srow, scol, erow, ecol, f'"{_BLANK}"')) + # A string that spans rows keeps all of them, so every later row stays where it was. + gap = "\n" * (erow - srow) + new = f'"""{_BLANK}{gap}"""' if gap else f'"{_BLANK}"' + edits.append((srow, scol, erow, ecol, new)) elif tok.type == _FSTRING_MIDDLE: # The f-string's literal text between `{...}` parts; FSTRING_START/END (the f"/" quote # delimiters) are left as-is, and any interpolated expression's own tokens (NAME, etc.) # still flow through the normal token handling above/below like any other reference. if tok.string: - edits.append((srow, scol, erow, ecol, _BLANK)) + edits.append((srow, scol, erow, ecol, _BLANK + "\n" * (erow - srow))) elif tok.type == tokenize.NUMBER: edits.append((srow, scol, erow, ecol, _BLANK)) elif tok.type == tokenize.COMMENT: @@ -225,7 +229,7 @@ def _names_of(tree, node_type, ranges) -> set[str]: # ── apply position edits to the source, preserving non-private lines verbatim ───────────── def _apply_edits(source: str, edits) -> str: - lines = source.splitlines(keepends=True) + lines = source_rows(source) # apply bottom-up so earlier edits don't shift later line indices / columns for srow, scol, erow, ecol, new in sorted(edits, reverse=True): if srow == erow: @@ -239,7 +243,7 @@ def _apply_edits(source: str, edits) -> str: # ── blank whole hidden lines, keeping indentation; leave blank lines untouched ──────────── def _apply_hides(text: str, hide_ranges) -> str: - lines = text.splitlines(keepends=True) + lines = source_rows(text) in_block = False # inside a run of consecutive hidden line numbers noted = False # has the explanatory note been added for the current block yet for i, line in enumerate(lines, 1): diff --git a/packages/syft-restrict/src/syft_restrict/runner.py b/packages/syft-restrict/src/syft_restrict/runner.py index 0c0b7d920c2..aa40244a692 100644 --- a/packages/syft-restrict/src/syft_restrict/runner.py +++ b/packages/syft-restrict/src/syft_restrict/runner.py @@ -11,9 +11,9 @@ from pydantic import BaseModel, Field -from .astutil import normalize_ranges, scan_file +from .astutil import normalize_ranges, scan_file, source_rows from .auditor import AuditReport, audit_allow_functions -from .errors import PolicyViolation, SourceEncodingError +from .errors import ObfuscationError, PolicyViolation, SourceEncodingError from .markers import parse_markers from .obfuscator import obfuscate as _obfuscate from .policy import Policy @@ -143,6 +143,7 @@ def _run( scan = scan_file(ast.parse(source), normalize_ranges(private)) obfuscated = _obfuscate(source, obfuscate_ranges, hide_ranges, scan) + _require_same_row_count(source, obfuscated) out_path = Path(out) if out is not None else path.with_suffix(".obfuscated.py") out_path.write_text(obfuscated, encoding="utf-8") @@ -165,6 +166,17 @@ def _run( ) +def _require_same_row_count(source: str, obfuscated: str) -> None: + """Raise ``ObfuscationError`` if ``obfuscated`` and ``source`` have different row counts.""" + # The hide ranges and the certificate's ranges count rows of the source. A copy with a + # different row count puts the hiding on the wrong rows. + n_src, n_out = len(source_rows(source)), len(source_rows(obfuscated)) + if n_out != n_src: + raise ObfuscationError( + f"the obfuscated copy has {n_out} rows and the source has {n_src}" + ) + + def _read_source(path: Path) -> str: """Read the file at ``path`` and return its text, with every line ending turned into ``\\n``. diff --git a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py index 3a659f54f0b..0870dbd1ab9 100644 --- a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py +++ b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py @@ -4,8 +4,9 @@ import shutil from pathlib import Path +import pytest from syft_restrict import obfuscate -from syft_restrict.astutil import scan_file +from syft_restrict.astutil import row_in_ranges, scan_file from syft_restrict.runner import _run FIXTURES = Path(__file__).parents[1] / "fixtures" @@ -103,6 +104,67 @@ def test_fstring_literal_text_is_blanked(): assert "■" in obf +def _assert_rows_aligned(source: str, obf: str, obfuscate_ranges, hide_ranges): + """Assert that public rows print unchanged and hidden code rows print as the hide marker.""" + src_rows, obf_rows = source.split("\n"), obf.split("\n") + assert len(obf_rows) == len(src_rows), obf + for row, (src_row, obf_row) in enumerate(zip(src_rows, obf_rows), 1): + if row_in_ranges(row, hide_ranges) and src_row.strip(): + assert obf_row.lstrip().startswith("■■■■■■■■"), (row, obf) + elif not row_in_ranges(row, [*obfuscate_ranges, *hide_ranges]): + assert obf_row == src_row, (row, obf) + + +@pytest.mark.parametrize( + "obfuscate_ranges, hide_ranges", + [([[3, 6]], []), ([[3, 3]], [[4, 6]])], + ids=["obfuscate", "obfuscate-then-hide"], +) +def test_formfeed_keeps_row_count(obfuscate_ranges, hide_ranges): + # A form feed must stay part of its row, the way tokenize counts rows. + source = ( + "import jax.numpy as jnp\n" + "# page\x0cbreak\n" + "def scale(x):\n" + " y = x * 2\n" + " z = y + 1\n" + " return z\n" + "\n" + "def main(x):\n" + " return scale(x)\n" + ) + private = [*obfuscate_ranges, *hide_ranges] + scan = scan_file(ast.parse(source), private) + obf = obfuscate(source, obfuscate_ranges, hide_ranges, scan) + _assert_rows_aligned(source, obf, obfuscate_ranges, hide_ranges) + + +@pytest.mark.parametrize( + "obfuscate_ranges, hide_ranges", + [([[3, 6]], []), ([[3, 6]], [[7, 7]])], + ids=["obfuscate", "obfuscate-then-hide"], +) +def test_multiline_docstring_keeps_row_count(obfuscate_ranges, hide_ranges): + # A string that spans several rows must keep all of them when it is blanked. + source = ( + "import jax.numpy as jnp\n" + "\n" + "class Model:\n" + ' """Line one.\n' + " Line two.\n" + ' Line three."""\n' + "SECRET = 1\n" + 'PUBLIC = "public"\n' + "\n" + "def main():\n" + " return PUBLIC\n" + ) + private = [*obfuscate_ranges, *hide_ranges] + scan = scan_file(ast.parse(source), private) + obf = obfuscate(source, obfuscate_ranges, hide_ranges, scan) + _assert_rows_aligned(source, obf, obfuscate_ranges, hide_ranges) + + def test_obfuscation_is_deterministic(tmp_path): dir_a = tmp_path / "a" dir_b = tmp_path / "b" diff --git a/packages/syft-restrict/tests/test_run.py b/packages/syft-restrict/tests/test_run.py index 54f28417183..0c1d34cf2ae 100644 --- a/packages/syft-restrict/tests/test_run.py +++ b/packages/syft-restrict/tests/test_run.py @@ -6,7 +6,14 @@ from pathlib import Path import pytest -from syft_restrict import MarkerError, PolicyViolation, SourceEncodingError, run +from syft_restrict import ( + MarkerError, + ObfuscationError, + PolicyViolation, + SourceEncodingError, + obfuscate, + run, +) from syft_restrict.runner import _run from verify.helpers import normalize_source @@ -195,6 +202,26 @@ def test__run_uses_explicit_ranges_and_ignores_markers(tmp_path): assert result.ok +def test_run_refuses_copy_with_different_row_count(tmp_path, monkeypatch): + # No known input makes obfuscate() change the row count, so the test drops a row itself. + src = tmp_path / "model.py" + shutil.copy(FIXTURES / "marked_model.py", src) + monkeypatch.setattr( + "syft_restrict.runner._obfuscate", + lambda *args: obfuscate(*args).split("\n", 1)[1], + ) + + # A broken copy is an obfuscator fault, not a verification result, so strict=False raises too. + with pytest.raises(ObfuscationError): + run( + src, + allow_functions=ALLOW_FUNCTIONS, + allow_operators=ALLOW_OPERATORS, + strict=False, + ) + assert not (tmp_path / "model.obfuscated.py").exists() + + def test_run_refuses_coding_cookie_that_changes_program(tmp_path): # The enclave loads model code with importlib, which honours a PEP 263 coding cookie. src = tmp_path / "model.py"