🗓 Jul 18, 2023 7:04:02 PM

🚀 return bytes instead of str in o/out 🧪 tests added/updated
securisec · Jul 18, 2023 · 74bf9d8 · 74bf9d8
1 parent 5ac00b5
commit 74bf9d8
Show file tree

Hide file tree

Showing 18 changed files with 237 additions and 234 deletions.
diff --git a/chepy/core.py b/chepy/core.py
@@ -638,6 +638,8 @@ def o(self):
         Returns:
             Any: Final output
         """
+        if isinstance(self.state, str):
+            return self.state.encode()
         return self.state
 
     @property
@@ -647,6 +649,8 @@ def out(self) -> Any:
         Returns:
             Any: Final output
         """
+        if isinstance(self.state, str):
+            return self.state.encode()
         return self.state
 
     @ChepyDecorators.call_stack

diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -55,11 +55,10 @@ def test_fire1():
 def test_fire2():
     assert (
         fire.Fire(Chepy, command=["abc", "-", "hmac_hash", "--digest", "md5"]).o
-        == "dd2701993d29fdd0b032c233cec63403"
+        == b"dd2701993d29fdd0b032c233cec63403"
     )
 
 
 def test_fire3():
     fire_obj = fire.Fire(Chepy, command=["abc", "-", "hmac_hash", "--digest", "md5"])
     assert type(fire_obj) == Chepy
-
diff --git a/tests/test_codetidy.py b/tests/test_codetidy.py
@@ -14,42 +14,44 @@ def test_beautify_json():
 
 
 def test_to_uppercase():
-    assert Chepy("some String").to_upper_case(by="word").o == "Some String"
-    assert Chepy("some String").to_upper_case(by="sentence").o == "Some string"
-    assert Chepy("some String").to_upper_case(by="all").o == "SOME STRING"
+    assert Chepy("some String").to_upper_case(by="word").o == b"Some String"
+    assert Chepy("some String").to_upper_case(by="sentence").o == b"Some string"
+    assert Chepy("some String").to_upper_case(by="all").o == b"SOME STRING"
 
 
 def test_to_snake_case():
-    assert Chepy("helloWorld").to_snake_case().o == "hello_world"
+    assert Chepy("helloWorld").to_snake_case().o == b"hello_world"
 
 
 def test_to_camel_case():
-    assert Chepy("some Data_test").to_camel_case().o == "someDataTest"
-    assert Chepy("some Data_test").to_camel_case(ignore_space=True).o == "some DataTest"
+    assert Chepy("some Data_test").to_camel_case().o == b"someDataTest"
+    assert (
+        Chepy("some Data_test").to_camel_case(ignore_space=True).o == b"some DataTest"
+    )
 
 
 def test_to_kebab_case():
-    assert Chepy("Some data_test").to_kebab_case().o == "some-data-test"
+    assert Chepy("Some data_test").to_kebab_case().o == b"some-data-test"
 
 
 def test_remove_whitespace():
     assert (
         Chepy("some    long space\n\ttab space\flol").remove_whitespace().o
-        == "somelongspacetabspacelol"
+        == b"somelongspacetabspacelol"
     )
 
 
 def test_swap_case():
-    assert Chepy("SoMe TeXt").swap_case().o == "sOmE tExT"
+    assert Chepy("SoMe TeXt").swap_case().o == b"sOmE tExT"
 
 
 def test_lower_case():
-    assert Chepy("HelLo WorLd").to_lower_case().o == "hello world"
+    assert Chepy("HelLo WorLd").to_lower_case().o == b"hello world"
 
 
 def test_leet_speak():
-    assert Chepy("somexValue").to_leetspeak().o == "50m3%V@1u3"
-    assert Chepy("somexValue").to_leetspeak(False).o == "50m3xVa1u3"
+    assert Chepy("somexValue").to_leetspeak().o == b"50m3%V@1u3"
+    assert Chepy("somexValue").to_leetspeak(False).o == b"50m3xVa1u3"
 
 
 def test_random_case():

diff --git a/tests/test_compression.py b/tests/test_compression.py
@@ -1,5 +1,6 @@
 from chepy import Chepy
 
+
 def test_fix_zip_header():
     assert (
         Chepy(
@@ -10,9 +11,10 @@ def test_fix_zip_header():
         .unzip_one("flag.txt")
         .trim()
         .o
-        == "LLS{tragic_number_more_like_magic_number}"
+        == b"LLS{tragic_number_more_like_magic_number}"
     )
 
+
 def test_zip_info():
     assert (
         Chepy("tests/files/test.zip")

diff --git a/tests/test_core.py b/tests/test_core.py
@@ -14,8 +14,8 @@ def test_states():
 
 
 def test_substring():
-    assert Chepy("some data").debug(True).substring("s(ome)", 1).o == "ome"
-    assert Chepy("some data").substring("s(ome)").o == "some"
+    assert Chepy("some data").debug(True).substring("s(ome)", 1).o == b"ome"
+    assert Chepy("some data").substring("s(ome)").o == b"some"
 
 
 def test_get_state():
@@ -35,7 +35,7 @@ def test_copy_state():
 
 
 def test_set_state():
-    assert Chepy("some data").set_state("new data").o == "new data"
+    assert Chepy("some data").set_state("new data").o == b"new data"
 
 
 def test_run_script():
@@ -48,8 +48,8 @@ def test_fork():
     assert Chepy("A", "B").fork(
         [("to_hex",), ("hmac_hash", {"key": "secret", "digest": "md5"})]
     ).states == {
-        0: "3e90033ea5422dafd81470dde4ffb37b",
-        1: "c474a4a957fe2018e2bffef53887ae22",
+        0: b"3e90033ea5422dafd81470dde4ffb37b",
+        1: b"c474a4a957fe2018e2bffef53887ae22",
     }
     assert c.fork([(c.to_hex,)]).states == {0: b"41", 1: b"42"}
 
@@ -144,7 +144,7 @@ def test_run_recipe():
             ]
         )
         .o
-        == "LOL\n"
+        == b"LOL\n"
     )
 
 
@@ -158,7 +158,7 @@ def test_recipe():
 
     assert (
         Chepy("tests/files/encoding").load_recipe(temp).o
-        == "StormCTF{Spot3:DcEC6181F48e3B9D3dF77Dd827BF34e0}"
+        == b"StormCTF{Spot3:DcEC6181F48e3B9D3dF77Dd827BF34e0}"
     )
     Path(temp).unlink()
 
@@ -171,25 +171,21 @@ def test_loop():
         == b"securisec"
     )
     c = Chepy("VmpGb2QxTXhXWGxTYmxKV1lrZDRWVmx0ZEV0alZsSllaVWRHYWxWVU1Eaz0=")
-    assert (
-        c.loop(6, c.from_base64)
-        .o
-        == b"securisec"
-    )
+    assert c.loop(6, c.from_base64).o == b"securisec"
 
 
 def test_loop_list():
     c = Chepy(["an", "array"])
     c.loop_list("to_hex").loop_list("hmac_hash", {"key": "secret"})
     assert c.o == [
-        "5cbe6ca2a66b380aec1449d4ebb0d40ac5e1b92e",
-        "30d75bf34740e8781cd4ec7b122e3efd8448e270",
+        b"5cbe6ca2a66b380aec1449d4ebb0d40ac5e1b92e",
+        b"30d75bf34740e8781cd4ec7b122e3efd8448e270",
     ]
     c1 = Chepy(["an", "array"])
     c1.loop_list("to_hex").loop_list(c.hmac_hash, {"key": "secret"})
     assert c1.o == [
-        "5cbe6ca2a66b380aec1449d4ebb0d40ac5e1b92e",
-        "30d75bf34740e8781cd4ec7b122e3efd8448e270",
+        b"5cbe6ca2a66b380aec1449d4ebb0d40ac5e1b92e",
+        b"30d75bf34740e8781cd4ec7b122e3efd8448e270",
     ]
 
 
@@ -198,19 +194,22 @@ def test_loop_dict():
     c = Chepy(data)
     c.loop_list("loop_dict", {"keys": ["some", "lol"], "callback": "to_upper_case"})
     assert c.o == [
-        {"some": "VAL"},
-        {"some": "ANOTHER"},
-        {"lol": "LOL"},
+        {"some": b"VAL"},
+        {"some": b"ANOTHER"},
+        {"lol": b"LOL"},
         {"another": "aaaa"},
     ]
 
-    d = Chepy({"some": "hahahaha", "lol": "aahahah"})
+    d = Chepy({"some": "hahahaha", "lol": b"aahahah"})
     d.loop_dict(["some"], d.to_upper_case)
-    assert d.o == {"some": "HAHAHAHA", "lol": "aahahah"}
+    assert d.o == {"some": b"HAHAHAHA", "lol": b"aahahah"}
 
     e = Chepy({"some": "hahahaha", "lol": "aahahah"})
     e.loop_dict(["some"], "hmac_hash", {"key": "secret"})
-    assert e.o == {"some": "99f77ec06a3c69a4a95371a7888245ba57f47f55", "lol": "aahahah"}
+    assert e.o == {
+        "some": b"99f77ec06a3c69a4a95371a7888245ba57f47f55",
+        "lol": "aahahah",
+    }
 
 
 def test_reset():