diff --git a/tod_attack_miner/cli.py b/tod_attack_miner/cli.py index 81fb4cf..55b8a90 100644 --- a/tod_attack_miner/cli.py +++ b/tod_attack_miner/cli.py @@ -30,6 +30,11 @@ def main(): default=None, help="If passed, filter TOD candidates that are {window-size} or more blocks apart", ) + parser.add_argument( + "--reset-db", + action="store_true", + help="Delete data from previous runs before starting to mine", + ) parser.add_argument("--postgres-user", type=str, default="postgres") parser.add_argument("--postgres-password", type=str, default="password") parser.add_argument("--postgres-host", type=str, default="localhost") @@ -49,6 +54,8 @@ def main(): if args.stats_only: print(json.dumps(miner.get_stats())) else: + if args.reset_db: + miner.reset_db() miner.fetch(int(args.from_block), int(args.to_block)) miner.compute_skelcodes() miner.find_collisions() diff --git a/tod_attack_miner/db/db.py b/tod_attack_miner/db/db.py index 46a4ca8..b10f91b 100644 --- a/tod_attack_miner/db/db.py +++ b/tod_attack_miner/db/db.py @@ -60,6 +60,25 @@ def _setup_tables(self): ) self._con.commit() + def reset(self): + with self._con.cursor() as cursor: + for name in _INDEXES: + cursor.execute( + cast( + psycopg.sql.SQL, + f"DROP INDEX IF EXISTS {name}", + ) + ) + for name in _TABLES: + cursor.execute( + cast( + psycopg.sql.SQL, + f"DROP TABLE IF EXISTS {name}", + ) + ) + self._con.commit() + self._setup_tables() + def insert_prestate(self, block_number: int, tx_index: int, prestate: TxPrestate): with self._con.cursor() as cursor: accesses: list[tuple[int, int, str, ACCESS_TYPE, str, str]] = [] diff --git a/tod_attack_miner/miner/miner.py b/tod_attack_miner/miner/miner.py index 6e07724..8ff552f 100644 --- a/tod_attack_miner/miner/miner.py +++ b/tod_attack_miner/miner/miner.py @@ -11,6 +11,9 @@ def __init__(self, rpc: RPC, db: DB) -> None: self._filter_stats = {"candidates": {}, "filtered": {}} self._original_collisions = {} + def reset_db(self) -> None: + self.db.reset() + def fetch(self, start: int, end: int) -> None: fetch_block_range(self.rpc, self.db, BlockRange(start, end))