🗓 Oct 22, 2023 11:45:00 PM

✨ new helper core methods
✨ flatten method
✨ search_list method
🧪 tests added/updated

chepy/core.py

@@ -543,7 +543,7 @@ def substring(self, pattern: str, group: int = 0):
 
     def _convert_to_bytes(self) -> bytes:
         """This method is used to coerce the current object in
-        the state variable into a string. The method should be
+        the state variable into a bytes. The method should be
         called inside any method that operates on a string object
         instead of calling `self.state` directly to avoid errors.
 
@@ -571,6 +571,35 @@ def _convert_to_bytes(self) -> bytes:
             # todo check more types here
             raise NotImplementedError
 
+    def _to_bytes(self, data: Any) -> bytes:  # pragma: no cover
+        """This method is used to coerce data to bytes. The method should be
+        called inside any method that operates on a string object
+        instead of calling `self.state` directly to avoid errors.
+
+        Raises:
+            NotImplementedError: If type coercian isn't available
+                for the current state type.
+        """
+        if isinstance(data, bytes):
+            return data
+        elif isinstance(data, str):
+            return data.encode()
+        elif isinstance(data, int):
+            return str(data).encode()
+        elif isinstance(data, dict):
+            return str(data).encode()
+        elif isinstance(data, list):
+            return str(data).encode()
+        elif isinstance(data, bool):
+            return str(data).encode()
+        elif isinstance(data, bytearray):
+            return bytes(data)
+        elif isinstance(data, float):
+            return bytearray(struct.pack("f", data))
+        else:  # pragma: no cover
+            # todo check more types here
+            raise NotImplementedError
+
     def _convert_to_bytearray(self) -> bytearray:
         """Attempts to coerce the current state into a
         `bytesarray` object

chepy/core.pyi

@@ -21,6 +21,7 @@ class ChepyCore:
     log_format: str = ...
     def __init__(self, *data: Any) -> None: ...
     def _convert_to_bytes(self) -> bytes: ...
+    def _to_bytes(self, data: Any) -> bytes: ...
     def _convert_to_bytearray(self) -> bytearray: ...
     def _convert_to_str(self) -> str: ...
     def _convert_to_int(self) -> int: ...

chepy/modules/dataformat.py

@@ -8,6 +8,7 @@
 import struct
 import pickle
 import string
+import itertools
 from random import randint
 from .internal.constants import Encoding
 from .internal.helpers import detect_delimiter, Rotate
@@ -1981,3 +1982,16 @@ def cut(self, start: int, end: int) -> DataFormatT:
         data = self._convert_to_bytes()
         self.state = data[start:end]
         return self
+
+    @ChepyDecorators.call_stack
+    def flatten(self) -> DataFormatT:
+        """Flatten a list of lists into a single list
+
+        Returns:
+            Chepy: The Chepy object.
+        """
+        try:
+            self.state = list(itertools.chain(*self.state))
+            return self
+        except:
+            return self

chepy/modules/dataformat.pyi

@@ -106,3 +106,4 @@ class DataFormat(ChepyCore):
     def to_base62(self: DataFormatT, alphabet: str="0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz") -> DataFormatT: ...
     def from_base62(self: DataFormatT, alphabet: str="0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz") -> DataFormatT: ...
     def cut(self: DataFormatT, start: int, end: int) -> DataFormatT: ...
+    def flatten(self: DataFormatT) -> DataFormatT: ...

chepy/modules/search.py

@@ -21,7 +21,7 @@ def search(self, pattern: Union[str, bytes]) -> SearchT:
         """Search. Group matches are returned as tuples.
 
         Args:
-            pattern (str): String pattern to search
+            pattern (Union[str, bytes]): Bytes pattern to search
 
         Returns:
             Chepy: The Chepy object.
@@ -34,6 +34,32 @@ def search(self, pattern: Union[str, bytes]) -> SearchT:
         self.state = re.findall(pattern, self._convert_to_bytes())
         return self
 
+    @ChepyDecorators.call_stack
+    def search_list(self, pattern: Union[str, bytes]) -> SearchT:
+        """Search all items in a list. List items are coerced into bytes first.
+        Group matches are returned as tuples.
+
+        Args:
+            pattern (Union[str, bytes]): Bytes pattern to search
+
+        Returns:
+            Chepy: The Chepy object.
+        """
+        assert isinstance(self.state, list), "State is not a list"
+
+        converted = [self._to_bytes(s) for s in self.state]
+        pattern = self._str_to_bytes(pattern)
+        pc = re.compile(pattern)
+
+        hold = []
+        for search in converted:
+            matches = pc.findall(search)
+            if len(matches) > 0:
+                hold.append(matches)
+
+        self.state = hold
+        return self
+
     @ChepyDecorators.call_stack
     def search_ctf_flags(self, prefix: str, postfix: str = ".+?\{*\}") -> SearchT:
         """Search CTF style flags.

chepy/modules/search.pyi

@@ -7,6 +7,7 @@ class Search(ChepyCore):
     def __init__(self, *data: Any) -> None: ...
     state: Any = ...
     def search(self: SearchT, pattern: Union[str, bytes]) -> SearchT: ...
+    def search_list(self: SearchT, pattern: Union[str, bytes]) -> SearchT: ...
     def search_ctf_flags(self: SearchT, prefix: str, postfix: str=...) -> SearchT: ...
     def search_slack_tokens(self: SearchT) -> SearchT: ...
     def search_slack_webhook(self: SearchT) -> SearchT: ...

tests/test_dataformat.py

@@ -637,3 +637,8 @@ def test_cut():
     assert Chepy("abcd").cut(1, 3).o == b"bc"
     assert Chepy({"a": 1}).cut(1, 3).o == b"'a"
     assert Chepy({"a": None}).cut(6, 9).o == b"Non"
+
+
+def test_flatten():
+    assert Chepy(["a", 1, True]).flatten().o == ["a", 1, True]
+    assert Chepy([["a"], [1], [True]]).flatten().o == ["a", 1, True]

tests/test_search.py

@@ -2,12 +2,18 @@
 
 
 def test_search():
-    assert Chepy("abcdefg123 and again abcdefg123").search(r"abc(de)fg(12)(3)").o == [
-        (b"de", b"12", b"3"),
+    assert Chepy("abcdefg123 and again abcdefg124").search(r"abc(de)fg(12)(\d)").o == [
         (b"de", b"12", b"3"),
+        (b"de", b"12", b"4"),
     ]
 
 
+def test_search_list():
+    assert Chepy(
+        ["InfoSeCon2023{1af5856c70878f8566085bc13849ef4d}", True, 123, ["a", "b"]]
+    ).search_list("Info.+").o == [[b"InfoSeCon2023{1af5856c70878f8566085bc13849ef4d}"]]
+
+
 def test_ctf_flags():
     assert (
         Chepy("tests/files/flags")