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"