Update solve.py

Author: skysoul1024

lab8/solve.py

@@ -1,28 +1,28 @@
-import angr
 import sys
-import claripy
+import angr
+
+def found_correct(state: angr.SimState):
+    return b"Correct!" in state.posix.dumps(1)
 
-angr.loggers.disable_root_logger()
+def avoid_wrong(state: angr.SimState):
+    return b"Wrong key!" in state.posix.dumps(1)
 
 def main():
-    proj = angr.Project("chal", auto_load_libs=False)
-    flag_bytes = [claripy.BVS(f'byte_{i}', 8) for i in range(8)]
-    flag = claripy.Concat(*flag_bytes)
+    angr.loggers.disable_root_logger()
 
-    state = proj.factory.full_init_state(
-        stdin = angr.SimFileStream(name='stdin', content=flag, has_end=True)
-    )
+    proj = angr.Project("./chal", auto_load_libs=False)
+    state = proj.factory.entry_state(stdin=angr.SimFile)
 
-    simgr = proj.factory.simulation_manager(state)
-    
-    simgr.explore(find=lambda s: b"flag" in s.posix.dumps(1))
+    simgr = proj.factory.simgr(state)
+    simgr.explore(find=found_correct, avoid=avoid_wrong)
 
-    if len(simgr.found) > 0:
+    if simgr.found:
         found_state = simgr.found[0]
-        solution = found_state.solver.eval(flag, cast_to=bytes)
+        solution = found_state.posix.dumps(0)
         sys.stdout.buffer.write(solution)
     else:
-        print("No solution found")
+        print("No solution found", file=sys.stderr)
+        exit(1)
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()