From c84f461b45d1e54324a0c5851442c22b04d0a72f Mon Sep 17 00:00:00 2001
From: Roman <zajic@zajic.net>
Date: Tue, 11 Feb 2025 14:20:58 +0800
Subject: [PATCH] fix: add remove_padding - parse exact result

---
 src/cli/nomos_cli.py                        | 27 ++++++++++++++++++---
 src/steps/da.py                             | 23 ++++++++++++++++++
 tests/data_integrity/test_data_integrity.py |  5 ++--
 3 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/src/cli/nomos_cli.py b/src/cli/nomos_cli.py
index b7449a6..2f83ead 100644
--- a/src/cli/nomos_cli.py
+++ b/src/cli/nomos_cli.py
@@ -1,4 +1,6 @@
+import json
 import os
+import re
 
 from src.data_storage import DS
 from src.libs.common import generate_log_prefix
@@ -8,6 +10,7 @@ from tenacity import retry, stop_after_delay, wait_fixed
 from src.cli.cli_vars import nomos_cli
 from src.docker_manager import DockerManager
 from src.env_vars import DOCKER_LOG_DIR, NOMOS_CLI
+from src.steps.da import remove_padding
 
 logger = get_custom_logger(__name__)
 
@@ -82,14 +85,30 @@ class NomosCli:
             self._container = None
             logger.debug("Container killed.")
 
-    def get_reconstruct_result(self):
+    def run_reconstruct(self, input_values=None, decode_only=True):
         keywords = ["Reconstructed data"]
+        self.run(input_values)
 
         log_stream = self._container.logs(stream=True)
 
         matches = self._docker_manager.search_log_for_keywords(self._log_path, keywords, False, log_stream)
-        # assert not matches, f"Reconstructed data not found {matches}"
-        for match in matches:
-            logger.debug(f"match {match}\n\n\n")
+        assert len(matches) > 0, f"Reconstructed data not found {matches}"
+
+        # Use regular expression that captures the byte list after "Reconstructed data"
+        result = re.sub(r".*Reconstructed data\s*(\[[^\]]+\]).*", r"\1", matches[keywords[0]][0])
+
+        result_bytes = []
+        try:
+            result_bytes = json.loads(result)
+        except Exception as ex:
+            logger.debug(f"Conversion to bytes failed with exception {ex}")
+
+        if decode_only:
+            result_bytes = result_bytes[:-31]
+
+        result_bytes = remove_padding(result_bytes)
+        result = bytes(result_bytes).decode("utf-8")
 
         DS.nomos_nodes.remove(self)
+
+        return result
diff --git a/src/steps/da.py b/src/steps/da.py
index 2422b48..a8eb854 100644
--- a/src/steps/da.py
+++ b/src/steps/da.py
@@ -24,6 +24,29 @@ def add_padding(orig_bytes):
     return padded_bytes
 
 
+def remove_padding(padded_bytes):
+    """
+    Removes PKCS#7-like padding from a list of bytes.
+    Raises:
+        ValueError: If the padding is incorrect.
+
+    Returns:
+        The original list of bytes without padding.
+    """
+    if not padded_bytes:
+        raise ValueError("The input is empty, cannot remove padding.")
+
+    padding_len = padded_bytes[-1]
+
+    if padding_len < 1 or padding_len > 31:
+        raise ValueError("Invalid padding length.")
+
+    if padded_bytes[-padding_len:] != [padding_len] * padding_len:
+        raise ValueError("Invalid padding bytes.")
+
+    return padded_bytes[:-padding_len]
+
+
 def prepare_dispersal_request(data, app_id, index):
     data_bytes = data.encode("utf-8")
     padded_bytes = add_padding(list(data_bytes))
diff --git a/tests/data_integrity/test_data_integrity.py b/tests/data_integrity/test_data_integrity.py
index 3f45ff0..1b7166a 100644
--- a/tests/data_integrity/test_data_integrity.py
+++ b/tests/data_integrity/test_data_integrity.py
@@ -35,7 +35,6 @@ class TestDataIntegrity(StepsDataAvailability):
         received_data = self.get_data_range(self.node2, [0] * 31 + [1], [0] * 8, [0] * 7 + [5])
         rcv_data_json = json.dumps(received_data)
         cli = NomosCli(command="reconstruct")
-        cli.run(input_values=[str(rcv_data_json)])
-        cli.get_reconstruct_result()
+        decoded_data = cli.run_reconstruct(input_values=[rcv_data_json])
 
-        # assert DATA_TO_DISPERSE[0] == bytes(received_data[0][1]).decode("utf-8")
+        assert DATA_TO_DISPERSE[0] == decoded_data, "Retrieved data are not same with original data"