From 95c7d04ada24821c14d935ecca7254a466ecd6ed Mon Sep 17 00:00:00 2001
From: Nikhil <nikhilparasaram@gmail.com>
Date: Tue, 22 Oct 2019 15:49:36 +0100
Subject: [PATCH] Add unit tests for keccak

---
 tests/laser/keccak_tests.py | 113 ++++++++++++++++++++++++++++++++++++
 1 file changed, 113 insertions(+)
 create mode 100644 tests/laser/keccak_tests.py

diff --git a/tests/laser/keccak_tests.py b/tests/laser/keccak_tests.py
new file mode 100644
index 00000000..d8f17eec
--- /dev/null
+++ b/tests/laser/keccak_tests.py
@@ -0,0 +1,113 @@
+from mythril.laser.smt import Solver, symbol_factory, And
+from mythril.laser.ethereum.keccak_function_manager import keccak_function_manager
+import z3
+import pytest
+
+
+@pytest.mark.parametrize(
+    "input1, input2, expected",
+    [
+        (symbol_factory.BitVecVal(100, 8), symbol_factory.BitVecVal(101, 8), z3.unsat),
+        (symbol_factory.BitVecVal(100, 8), symbol_factory.BitVecVal(100, 16), z3.unsat),
+        (symbol_factory.BitVecVal(100, 8), symbol_factory.BitVecVal(100, 8), z3.sat),
+        (
+            symbol_factory.BitVecSym("N1", 256),
+            symbol_factory.BitVecSym("N2", 256),
+            z3.sat,
+        ),
+    ],
+)
+def test_keccak_basic(input1, input2, expected):
+    s = Solver()
+
+    o1, c1 = keccak_function_manager.create_keccak(input1)
+    o2, c2 = keccak_function_manager.create_keccak(input2)
+    s.add(c1)
+    s.add(c2)
+
+    s.add(o1 == o2)
+    assert s.check() == expected
+
+
+def test_keccak_complex_eq():
+    """
+    check for keccak(keccak(b)*2) == keccak(keccak(a)*2) && a != b
+    :return:
+    """
+    s = Solver()
+    a = symbol_factory.BitVecSym("a", 160)
+    b = symbol_factory.BitVecSym("b", 160)
+    o1, c1 = keccak_function_manager.create_keccak(a)
+    o2, c2 = keccak_function_manager.create_keccak(b)
+    s.add(And(c1, c2))
+    two = symbol_factory.BitVecVal(2, 256)
+    o1 = two * o1
+    o2 = two * o2
+    o1, c1 = keccak_function_manager.create_keccak(o1)
+    o2, c2 = keccak_function_manager.create_keccak(o2)
+
+    s.add(And(c1, c2))
+    s.add(o1 == o2)
+    s.add(a != b)
+
+    assert s.check() == z3.unsat
+
+
+def test_keccak_complex_eq2():
+    """
+    check for keccak(keccak(b)*2) == keccak(keccak(a)*2)
+    This isn't combined with prev test because incremental solving here requires extra-extra work
+    (solution is literally the opposite of prev one) so it will take forever to solve.
+    :return:
+    """
+    s = Solver()
+    a = symbol_factory.BitVecSym("a", 160)
+    b = symbol_factory.BitVecSym("b", 160)
+    o1, c1 = keccak_function_manager.create_keccak(a)
+    o2, c2 = keccak_function_manager.create_keccak(b)
+    s.add(And(c1, c2))
+    two = symbol_factory.BitVecVal(2, 256)
+    o1 = two * o1
+    o2 = two * o2
+    o1, c1 = keccak_function_manager.create_keccak(o1)
+    o2, c2 = keccak_function_manager.create_keccak(o2)
+
+    s.add(And(c1, c2))
+    s.add(o1 == o2)
+
+    assert s.check() == z3.sat
+
+
+def test_keccak_simple_number():
+    """
+    check for keccak(b) == 10
+    :return:
+    """
+    s = Solver()
+    a = symbol_factory.BitVecSym("a", 160)
+    ten = symbol_factory.BitVecVal(10, 256)
+    o, c = keccak_function_manager.create_keccak(a)
+
+    s.add(c)
+    s.add(ten == o)
+
+    assert s.check() == z3.unsat
+
+
+def test_keccak_other_num():
+    """
+    check keccak(keccak(a)*2) == b
+    :return:
+    """
+    s = Solver()
+    a = symbol_factory.BitVecSym("a", 160)
+    b = symbol_factory.BitVecSym("b", 256)
+    o, c = keccak_function_manager.create_keccak(a)
+    two = symbol_factory.BitVecVal(2, 256)
+    o = two * o
+    s.add(c)
+    o, c = keccak_function_manager.create_keccak(o)
+    s.add(c)
+    s.add(b == o)
+
+    assert s.check() == z3.sat