Skip to content

Instantly share code, notes, and snippets.

@L3viathan
Created August 13, 2024 12:09
Show Gist options
  • Save L3viathan/3e15d5db733683d3eb74282972ba3581 to your computer and use it in GitHub Desktop.
Save L3viathan/3e15d5db733683d3eb74282972ba3581 to your computer and use it in GitHub Desktop.
Protect parts of your code with a contract. Contract-protected code can't be (accidentally) changed.
import inspect
import ast
import hashlib
import sys
import textwrap
from contextlib import contextmanager
def indent(line):
return len(line) - len(line.lstrip(" "))
@contextmanager
def contract(description, *, signature=None):
f = sys._getframe(2)
lines = inspect.getsource(f).split("\n")
start = f.f_lineno
start_indent = indent(lines[start])
stop = next(
i for i in range(start + 1, len(lines)) if indent(lines[i]) < start_indent
)
content = textwrap.dedent("\n".join(lines[start:stop]))
hash_ = hashlib.sha1(ast.unparse(ast.parse(content)).encode("utf-8")).hexdigest()[:8]
if signature is None:
raise ValueError(f"No signature passed, should be: {hash_}")
if signature != hash_:
raise RuntimeError(f"Contract-protected code was modified. The following contract needs to be fulfilled by the code:\n\n{description}\n\nRevert your changes or adjust the signature (by leaving it empty for one run), if you have documented your change.")
yield
if __name__ == "__main__":
with contract("Prints hello and bye", signature="faa7b015"):
print("hello")
if True:
print("bye")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment