From e5697e6e97227b3b2c9bab808feecbef1723522a Mon Sep 17 00:00:00 2001 From: Aashutosh Poudel Date: Thu, 5 Nov 2026 20:59:31 +0000 Subject: [PATCH 1/3] fix(restrict): count rows the way tokenize does, and refuse a misaligned copy The obfuscator computed its edits and the certificate ranges from tokenize rows, then applied them to rows from str.splitlines. splitlines also breaks on a form feed and seven other characters that tokenize treats as ordinary text. One form feed before a private block moved every edit up a row: the private class was printed in the clear and a public comment was rewritten (finding A3, form feed). A new astutil.source_rows() splits on \n only, like tokenize. The obfuscator's edit and hide steps and the marker parser now use it. run() now compares the row counts of the obfuscated copy and the source. On a mismatch it writes no copy and issues no certificate. It reports a new row-count-mismatch violation (raised as PolicyViolation under strict). This also refuses the docstring-shift form of A3, where a multi-line string in an obfuscate block collapses to one row, until the obfuscator keeps those rows. Tests: a form feed inside a public comment before an obfuscate block, and before an obfuscate-then-hide block, keeps the row count. Both fail on the old code. Co-Authored-By: Claude Opus 5.5 --- packages/syft-restrict/docs/blacklist.md | 35 ++++++++++--------- .../src/syft_restrict/astutil.py | 14 ++++++++ .../src/syft_restrict/markers.py | 3 +- .../src/syft_restrict/obfuscator.py | 5 +-- .../syft-restrict/src/syft_restrict/runner.py | 14 +++++++- .../src/syft_restrict/verifier.py | 1 + .../tests/obfuscate/test_obfuscate.py | 26 ++++++++++++++ 7 files changed, 77 insertions(+), 21 deletions(-) diff --git a/packages/syft-restrict/docs/blacklist.md b/packages/syft-restrict/docs/blacklist.md index 89cb164ca86..e7a9df80b7a 100644 --- a/packages/syft-restrict/docs/blacklist.md +++ b/packages/syft-restrict/docs/blacklist.md @@ -12,23 +12,24 @@ This document lists everything the checker rejects in the **private** region. ## Index of violation codes -| Code | Meaning | -| ------------------- | ------------------------------------------------------------------------ | -| `node-type` | Syntax not on the allow list (walrus, `match`, …) | -| `banned-name` | Reference or call to a banned builtin (`open`, `eval`, …) | -| `banned-construct` | Forbidden statement/expression form (import, `with`, f-string, …) | -| `call-not-allowed` | Library path not allow-listed (or hit by `disallow_functions`) | -| `call-unresolved` | Call target not provably safe | -| `method-on-value` | Named method on an unknown value (`x.reshape`) | -| `attr-on-value` | Attribute read/write on an unknown value, or bad `self` chain | -| `attr-not-allowed` | Library attribute path not allow-listed | -| `dunder-attr` | Dunder attribute (`.__class__`, …) | -| `dunder-name` | Bare dunder name (`__class__`) | -| `dunder-def` | Defining a forbidden magic method | -| `class-keyword` | e.g. `metaclass=` | -| `class-base` | Base class not allow-listed | -| `reserved-name` | Rebinding a trusted name (`jnp`, public wrapper, private def, `self`, …) | -| `operator-disabled` | Operator used without enabling its group in `allow_operators` | +| Code | Meaning | +| -------------------- | ------------------------------------------------------------------------ | +| `node-type` | Syntax not on the allow list (walrus, `match`, …) | +| `banned-name` | Reference or call to a banned builtin (`open`, `eval`, …) | +| `banned-construct` | Forbidden statement/expression form (import, `with`, f-string, …) | +| `call-not-allowed` | Library path not allow-listed (or hit by `disallow_functions`) | +| `call-unresolved` | Call target not provably safe | +| `method-on-value` | Named method on an unknown value (`x.reshape`) | +| `attr-on-value` | Attribute read/write on an unknown value, or bad `self` chain | +| `attr-not-allowed` | Library attribute path not allow-listed | +| `dunder-attr` | Dunder attribute (`.__class__`, …) | +| `dunder-name` | Bare dunder name (`__class__`) | +| `dunder-def` | Defining a forbidden magic method | +| `class-keyword` | e.g. `metaclass=` | +| `class-base` | Base class not allow-listed | +| `reserved-name` | Rebinding a trusted name (`jnp`, public wrapper, private def, `self`, …) | +| `operator-disabled` | Operator used without enabling its group in `allow_operators` | +| `row-count-mismatch` | The obfuscated copy has a different row count than the source | --- diff --git a/packages/syft-restrict/src/syft_restrict/astutil.py b/packages/syft-restrict/src/syft_restrict/astutil.py index 3e0ca3e790b..51c9b2480d2 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,19 @@ def normalize_ranges(private) -> list[tuple[int, int]]: return ranges +def source_rows(text: str) -> list[str]: + """Split ``text`` into rows on newline characters, ``\\n``, similar to + ``tokenize``. + + This function replaces ``str.splitlines`` which splits on form feeds, + file separators, and other separator characters. + + See: https://docs.python.org/3/builtins/stdtypes.html#str.splitlines + """ + + 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/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..8e815aeeba0 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 @@ -225,7 +226,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 +240,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 b7c51cec0ff..e14b4407979 100644 --- a/packages/syft-restrict/src/syft_restrict/runner.py +++ b/packages/syft-restrict/src/syft_restrict/runner.py @@ -8,7 +8,7 @@ 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 from .markers import parse_markers @@ -140,6 +140,18 @@ def _run( scan = scan_file(ast.parse(source), normalize_ranges(private)) obfuscated = _obfuscate(source, obfuscate_ranges, hide_ranges, scan) + # The certificate's ranges count rows of the source. A copy with a different row count would + # put the hiding on the wrong rows, so refuse to write it or certify it. + n_src, n_out = len(source_rows(source)), len(source_rows(obfuscated)) + if n_out != n_src: + mismatch = Violation( + line=min((lo for lo, _ in normalize_ranges(private)), default=1), + code="row-count-mismatch", + message=f"the obfuscated copy has {n_out} rows and the source has {n_src}", + ) + if strict: + raise PolicyViolation([mismatch]) + return RunResult(ok=False, violations=[mismatch], audit=audit) out_path = Path(out) if out is not None else path.with_suffix(".obfuscated.py") out_path.write_text(obfuscated) diff --git a/packages/syft-restrict/src/syft_restrict/verifier.py b/packages/syft-restrict/src/syft_restrict/verifier.py index c2cc63548bc..184fa257aca 100644 --- a/packages/syft-restrict/src/syft_restrict/verifier.py +++ b/packages/syft-restrict/src/syft_restrict/verifier.py @@ -153,6 +153,7 @@ "operator-disabled", # _require_bundle "duplicate-method", # _forbid_duplicate_methods "star-import", # visit (`from ... import *` anywhere in the file) + "row-count-mismatch", # run (display copy row count differs from the source) ] diff --git a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py index 3a659f54f0b..23839ffb7c0 100644 --- a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py +++ b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py @@ -4,6 +4,8 @@ import shutil from pathlib import Path +import pytest + from syft_restrict import obfuscate from syft_restrict.astutil import scan_file from syft_restrict.runner import _run @@ -103,6 +105,30 @@ def test_fstring_literal_text_is_blanked(): assert "■" in 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 obf.count("\n") == source.count("\n"), obf + + def test_obfuscation_is_deterministic(tmp_path): dir_a = tmp_path / "a" dir_b = tmp_path / "b" From 635229dd17a8c1caa714ee0a0009b676dd97c093 Mon Sep 17 00:00:00 2001 From: Aashutosh Poudel Date: Thu, 5 Nov 2026 21:42:32 +0000 Subject: [PATCH 2/3] fix(restrict): keep every row of a multi-line string when it is blanked MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The obfuscator replaced a string token with one "■", however many rows it spanned. A five-line docstring in an obfuscate block collapsed to one row. The hide step then used source row numbers on the shorter text: it blanked a public line and printed the hidden lines in the clear (finding A3, docstring shift). Since the previous commit, run() refused such a file. A string that spans rows is now replaced by a triple-quoted "■" with the same number of newlines. Multi-line f-string text keeps its newlines too. The copy keeps the source's row count, so run() accepts the file again and the hide step blanks the right rows. Tests: a three-line docstring in an obfuscate block, with and without a hide block after it, keeps the row count. Co-Authored-By: Claude Opus 5.5 --- .../src/syft_restrict/obfuscator.py | 7 +++-- .../tests/obfuscate/test_obfuscate.py | 26 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/packages/syft-restrict/src/syft_restrict/obfuscator.py b/packages/syft-restrict/src/syft_restrict/obfuscator.py index 8e815aeeba0..732975d7cfe 100644 --- a/packages/syft-restrict/src/syft_restrict/obfuscator.py +++ b/packages/syft-restrict/src/syft_restrict/obfuscator.py @@ -113,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: diff --git a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py index 23839ffb7c0..9048b6ff3cc 100644 --- a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py +++ b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py @@ -129,6 +129,32 @@ def test_formfeed_keeps_row_count(obfuscate_ranges, hide_ranges): assert obf.count("\n") == source.count("\n"), obf +@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 obf.count("\n") == source.count("\n"), obf + + def test_obfuscation_is_deterministic(tmp_path): dir_a = tmp_path / "a" dir_b = tmp_path / "b" From 6ab281aec84bb53ffa767d043417cf9beef04a8c Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Fri, 13 Nov 2026 16:46:32 -0300 Subject: [PATCH 3/3] fix: raise ObfuscationError if the display copy loses or adds rows A different row count in the display copy is an error in the obfuscator. It is not a problem in the private code. Therefore run() now raises ObfuscationError for it, also with strict=False. The row-count-mismatch violation code no longer exists. The check is in _require_same_row_count(). A new test makes the copy lose a row and expects the error. The row-count tests of the obfuscator also check that hidden rows show the hide marker and that public rows do not change. --- packages/syft-restrict/docs/blacklist.md | 35 +++++++++---------- .../src/syft_restrict/__init__.py | 9 ++++- .../src/syft_restrict/astutil.py | 10 ++---- .../syft-restrict/src/syft_restrict/errors.py | 4 +++ .../syft-restrict/src/syft_restrict/runner.py | 26 +++++++------- .../src/syft_restrict/verifier.py | 1 - .../tests/obfuscate/test_obfuscate.py | 18 +++++++--- packages/syft-restrict/tests/test_run.py | 29 ++++++++++++++- 8 files changed, 87 insertions(+), 45 deletions(-) diff --git a/packages/syft-restrict/docs/blacklist.md b/packages/syft-restrict/docs/blacklist.md index e7a9df80b7a..89cb164ca86 100644 --- a/packages/syft-restrict/docs/blacklist.md +++ b/packages/syft-restrict/docs/blacklist.md @@ -12,24 +12,23 @@ This document lists everything the checker rejects in the **private** region. ## Index of violation codes -| Code | Meaning | -| -------------------- | ------------------------------------------------------------------------ | -| `node-type` | Syntax not on the allow list (walrus, `match`, …) | -| `banned-name` | Reference or call to a banned builtin (`open`, `eval`, …) | -| `banned-construct` | Forbidden statement/expression form (import, `with`, f-string, …) | -| `call-not-allowed` | Library path not allow-listed (or hit by `disallow_functions`) | -| `call-unresolved` | Call target not provably safe | -| `method-on-value` | Named method on an unknown value (`x.reshape`) | -| `attr-on-value` | Attribute read/write on an unknown value, or bad `self` chain | -| `attr-not-allowed` | Library attribute path not allow-listed | -| `dunder-attr` | Dunder attribute (`.__class__`, …) | -| `dunder-name` | Bare dunder name (`__class__`) | -| `dunder-def` | Defining a forbidden magic method | -| `class-keyword` | e.g. `metaclass=` | -| `class-base` | Base class not allow-listed | -| `reserved-name` | Rebinding a trusted name (`jnp`, public wrapper, private def, `self`, …) | -| `operator-disabled` | Operator used without enabling its group in `allow_operators` | -| `row-count-mismatch` | The obfuscated copy has a different row count than the source | +| Code | Meaning | +| ------------------- | ------------------------------------------------------------------------ | +| `node-type` | Syntax not on the allow list (walrus, `match`, …) | +| `banned-name` | Reference or call to a banned builtin (`open`, `eval`, …) | +| `banned-construct` | Forbidden statement/expression form (import, `with`, f-string, …) | +| `call-not-allowed` | Library path not allow-listed (or hit by `disallow_functions`) | +| `call-unresolved` | Call target not provably safe | +| `method-on-value` | Named method on an unknown value (`x.reshape`) | +| `attr-on-value` | Attribute read/write on an unknown value, or bad `self` chain | +| `attr-not-allowed` | Library attribute path not allow-listed | +| `dunder-attr` | Dunder attribute (`.__class__`, …) | +| `dunder-name` | Bare dunder name (`__class__`) | +| `dunder-def` | Defining a forbidden magic method | +| `class-keyword` | e.g. `metaclass=` | +| `class-base` | Base class not allow-listed | +| `reserved-name` | Rebinding a trusted name (`jnp`, public wrapper, private def, `self`, …) | +| `operator-disabled` | Operator used without enabling its group in `allow_operators` | --- 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 51c9b2480d2..75269d20756 100644 --- a/packages/syft-restrict/src/syft_restrict/astutil.py +++ b/packages/syft-restrict/src/syft_restrict/astutil.py @@ -24,15 +24,11 @@ def normalize_ranges(private) -> list[tuple[int, int]]: def source_rows(text: str) -> list[str]: - """Split ``text`` into rows on newline characters, ``\\n``, similar to - ``tokenize``. + """Split ``text`` into rows on ``\\n`` only, as ``tokenize`` counts them. - This function replaces ``str.splitlines`` which splits on form feeds, - file separators, and other separator characters. - - See: https://docs.python.org/3/builtins/stdtypes.html#str.splitlines + 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() 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/runner.py b/packages/syft-restrict/src/syft_restrict/runner.py index 632a867c4db..aa40244a692 100644 --- a/packages/syft-restrict/src/syft_restrict/runner.py +++ b/packages/syft-restrict/src/syft_restrict/runner.py @@ -13,7 +13,7 @@ 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,18 +143,7 @@ def _run( scan = scan_file(ast.parse(source), normalize_ranges(private)) obfuscated = _obfuscate(source, obfuscate_ranges, hide_ranges, scan) - # The certificate's ranges count rows of the source. A copy with a different row count would - # put the hiding on the wrong rows, so refuse to write it or certify it. - n_src, n_out = len(source_rows(source)), len(source_rows(obfuscated)) - if n_out != n_src: - mismatch = Violation( - line=min((lo for lo, _ in normalize_ranges(private)), default=1), - code="row-count-mismatch", - message=f"the obfuscated copy has {n_out} rows and the source has {n_src}", - ) - if strict: - raise PolicyViolation([mismatch]) - return RunResult(ok=False, violations=[mismatch], audit=audit) + _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") @@ -177,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/src/syft_restrict/verifier.py b/packages/syft-restrict/src/syft_restrict/verifier.py index 184fa257aca..c2cc63548bc 100644 --- a/packages/syft-restrict/src/syft_restrict/verifier.py +++ b/packages/syft-restrict/src/syft_restrict/verifier.py @@ -153,7 +153,6 @@ "operator-disabled", # _require_bundle "duplicate-method", # _forbid_duplicate_methods "star-import", # visit (`from ... import *` anywhere in the file) - "row-count-mismatch", # run (display copy row count differs from the source) ] diff --git a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py index 9048b6ff3cc..0870dbd1ab9 100644 --- a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py +++ b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py @@ -5,9 +5,8 @@ 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" @@ -105,6 +104,17 @@ 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]])], @@ -126,7 +136,7 @@ def test_formfeed_keeps_row_count(obfuscate_ranges, hide_ranges): private = [*obfuscate_ranges, *hide_ranges] scan = scan_file(ast.parse(source), private) obf = obfuscate(source, obfuscate_ranges, hide_ranges, scan) - assert obf.count("\n") == source.count("\n"), obf + _assert_rows_aligned(source, obf, obfuscate_ranges, hide_ranges) @pytest.mark.parametrize( @@ -152,7 +162,7 @@ def test_multiline_docstring_keeps_row_count(obfuscate_ranges, hide_ranges): private = [*obfuscate_ranges, *hide_ranges] scan = scan_file(ast.parse(source), private) obf = obfuscate(source, obfuscate_ranges, hide_ranges, scan) - assert obf.count("\n") == source.count("\n"), obf + _assert_rows_aligned(source, obf, obfuscate_ranges, hide_ranges) def test_obfuscation_is_deterministic(tmp_path): 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"