From e3cbd598ec645e5b8875c48c5683e5c0d24e16af Mon Sep 17 00:00:00 2001
From: Pierre Lalet <pierre@droids-corp.org>
Date: Wed, 20 Dec 2023 22:37:26 +0100
Subject: [PATCH] CLI: skip already used items in `hf mf elog --decrypt`

This (often largely) improves the speed of the decrypt process. On my
laptop, with the same logs (37 records for one block and 37 records
for another block), here are the performances, as measuerd using a
simple command:

```bash
time echo -e "hw connect\nhf mf elog --decrypt\nhw disconnect" | ./chameleon_cli_main.py
```

- Before parallelisation (#187): 14m59,277s
- With parallelisation (current main): 6m13,513s
- With item skipping (this PR): 2m42,491s
---
 CHANGELOG.md                          |  1 +
 software/script/chameleon_cli_unit.py | 69 +++++++++++++++++++++------
 2 files changed, 55 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1b1fc399..6781907a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@ All notable changes to this project will be documented in this file.
 This project uses the changelog in accordance with [keepchangelog](http://keepachangelog.com/). Please use this to write notable changes, which is not the same as git commit log...
 
 ## [unreleased][unreleased]
+ - Skip already used items `hf mf elog --decrypt` (@p-l-)
  - Parallelize mfkey32v2 processes called from CLI (@p-l-)
  - Added support for mifare classic value block operations (@taichunmin)
  - Added regression tests (@doegox)
diff --git a/software/script/chameleon_cli_unit.py b/software/script/chameleon_cli_unit.py
index c47487b7..c42b80d1 100644
--- a/software/script/chameleon_cli_unit.py
+++ b/software/script/chameleon_cli_unit.py
@@ -1044,10 +1044,57 @@ def _run_mfkey32v2(items):
     ).stdout
     sea_obj = _KEY.search(output_str)
     if sea_obj is not None:
-        return sea_obj[0]
+        return sea_obj[0], items
     return None
 
 
+class ItemGenerator:
+    def __init__(self, rs, i=0, j=1):
+        self.rs = rs
+        self.i = 0
+        self.j = 1
+        self.found = set()
+        self.keys = set()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            item_i = self.rs[self.i]
+        except IndexError:
+            raise StopIteration
+        if self.key_from_item(item_i) in self.found:
+            self.i += 1
+            self.j = self.i + 1
+            return next(self)
+        try:
+            item_j = self.rs[self.j]
+        except IndexError:
+            self.i += 1
+            self.j = self.i + 1
+            return next(self)
+        self.j += 1
+        if self.key_from_item(item_j) in self.found:
+            return next(self)
+        return item_i, item_j
+
+    @staticmethod
+    def key_from_item(item):
+        return "{uid}-{nt}-{nr}-{ar}".format(**item)
+
+    def key_found(self, key, items):
+        self.keys.add(key)
+        for item in items:
+            try:
+                if item == self.rs[self.i]:
+                    self.i += 1
+                    self.j = self.i + 1
+            except IndexError:
+                break
+        self.found.update(self.key_from_item(item) for item in items)
+
+
 @hf_mf.command('elog')
 class HFMFELog(DeviceRequiredUnit):
     detection_log_size = 18
@@ -1069,24 +1116,16 @@ def decrypt_by_list(self, rs: list):
         msg2 = f"/{(len(rs)*(len(rs)-1))//2} combinations. "
         msg3 = " key(s) found"
         n = 1
-        keys = set()
+        gen = ItemGenerator(rs)
         with Pool(cpu_count()) as pool:
-            for key in pool.imap(
-                _run_mfkey32v2,
-                (
-                    (item0, rs[j])
-                    for i, item0 in enumerate(rs)
-                    for j in range(i + 1, len(rs))
-                ),
-            ):
+            for result in pool.imap(_run_mfkey32v2, gen):
                 # TODO: if some keys already recovered, test them on item before running mfkey32 on item
-                # TODO: if some keys already recovered, remove corresponding items
-                if key is not None:
-                    keys.add(key)
-                print(f"{msg1}{n}{msg2}{len(keys)}{msg3}\r", end="")
+                if result is not None:
+                    gen.key_found(*result)
+                print(f"{msg1}{n}{msg2}{len(gen.keys)}{msg3}\r", end="")
                 n += 1
         print()
-        return keys
+        return gen.keys
 
     def on_exec(self, args: argparse.Namespace):
         if not args.decrypt: