From 585881265346a3150ed621cbbb297f202e816ec9 Mon Sep 17 00:00:00 2001
From: Guillaume Pujol <guill.p.linux@gmail.com>
Date: Sat, 5 Aug 2023 21:27:15 +0200
Subject: [PATCH] added support for Caesar cipher (#4)

* added support for Caesar cipher
---
 binapy/encoding/__init__.py |  1 +
 binapy/encoding/dumb.py     | 69 +++++++++++++++++++++++++++++++++++++
 tests/test_encoding.py      | 19 ++++++++++
 3 files changed, 89 insertions(+)
 create mode 100644 binapy/encoding/dumb.py

diff --git a/binapy/encoding/__init__.py b/binapy/encoding/__init__.py
index 90191b2..6968612 100644
--- a/binapy/encoding/__init__.py
+++ b/binapy/encoding/__init__.py
@@ -4,5 +4,6 @@
 """
 
 from .base64 import *
+from .dumb import *
 from .hex import *
 from .url import *
diff --git a/binapy/encoding/dumb.py b/binapy/encoding/dumb.py
new file mode 100644
index 0000000..a29e84c
--- /dev/null
+++ b/binapy/encoding/dumb.py
@@ -0,0 +1,69 @@
+"""Implement support for 'dumb' ciphers such as Caesar cipher."""
+
+import string
+from typing import Union
+
+from binapy import binapy_decoder, binapy_encoder
+
+
+@binapy_encoder("caesar")
+def encode_caesar(
+    bp: bytes,
+    shift: int,
+    alphabet: Union[None, str, bytes] = None,
+) -> bytes:
+    """Encode data with Caesar cipher.
+
+    This shifts each character from `bp` by `shift` positions in the given `alphabet`.
+    Characters from `bp` that are not in the alphabet are left as-is.
+    Alphabet is usually `string.ascii_lowercase`, `string.ascii_uppercase`, but you may pass any
+    alphabet, either as a `str` (which will be encoded using 'utf-8'), or as `bytes` directly.
+    By default, alphabet will be auto-detected:
+
+    - `string.ascii_uppercase` if all character from the input are uppercase letters ASCII codes
+    - `string.ascii_lowercase` if all character from the input are lowercase letters ASCII codes
+    - `string.ascii_letters` if all character from the input are letters (both upper and lower case) ASCII codes
+    - the full ASCII range (0-127) if all characters are valid ASCII
+    - the full octect range (0-255) otherwise
+
+    Args:
+        bp: input data.
+        shift: number of places to shift each character in the alphabet.
+        alphabet: alphabet to use. Leave `None` to try to auto-detect alphabet
+
+    Returns:
+        the result of applying Caesar-cipher with `shift` positions to `bp`.
+    """
+    if not alphabet:
+        if all(65 <= c <= 90 for c in bp):
+            alphabet = string.ascii_uppercase
+        elif all(97 <= c <= 122 for c in bp):
+            alphabet = string.ascii_lowercase
+        elif all(65 <= c <= 90 or 97 <= c <= 122 for c in bp):
+            alphabet = string.ascii_letters
+        elif all(0 <= c <= 127 for c in bp):
+            alphabet = bytes(range(128))
+        else:
+            alphabet = bytes(range(256))
+
+    if isinstance(alphabet, str):
+        alphabet = alphabet.encode()
+
+    return bytes(
+        alphabet[(alphabet.index(c) + shift) % len(alphabet)] if c in alphabet else c
+        for c in bp
+    )
+
+
+@binapy_decoder("caesar")
+def decode_caesar(
+    bp: bytes,
+    shift: int,
+    alphabet: Union[None, str, bytes] = None,
+) -> bytes:
+    """Decode data with Caesar cipher.
+
+    Since encoding and decoding are symmetric, this is just an alias to `encode_caesar()` with an
+    opposite shift.
+    """
+    return encode_caesar(bp, -shift, alphabet)
diff --git a/tests/test_encoding.py b/tests/test_encoding.py
index 6319e8d..e93f20f 100644
--- a/tests/test_encoding.py
+++ b/tests/test_encoding.py
@@ -1,3 +1,5 @@
+import string
+
 import pytest
 
 from binapy import BinaPy
@@ -37,3 +39,20 @@ def test_url() -> None:
         bp.encode_to("url", plus_spaces=False)
         == b"https%3A//localhost%3A3200/foo%3Fbar%3Dab%20cd"
     )
+
+
+def test_caesar() -> None:
+    assert BinaPy("caesar13").to("caesar", 13, string.ascii_lowercase) == b"pnrfne13"
+    assert (
+        BinaPy(string.ascii_lowercase).to("caesar", 13, string.ascii_lowercase)
+        == b"nopqrstuvwxyzabcdefghijklm"
+    )
+    assert (
+        BinaPy(string.ascii_letters).to("caesar", 13, string.ascii_letters)
+        == b"nopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklm"
+    )
+
+    for data in (b"caesar", b"CAESAR", b"FooBAR", bytes(range(128))):
+        assert BinaPy(data).to("caesar", 4).decode_from("caesar", 4) == data
+
+    assert BinaPy(b"\x00\xff").to("caesar", 1) == b"\x01\x00"