6
\$\begingroup\$

I made this code to encrypt a string with the One Time Pad method. Is this pythonic code, and are there obvious ways I can improve?
The program accepts a plaintext and can either accept an inputted key or generate one itself. It adds the key and plaintext character by character, keeping it within the bounds of the string. e.g. "This is an example" encoded with "Mary had a little lamb" gives "8rJ!~zCc9x~zPDPKzd".

import string
import secrets

random_generator = secrets.SystemRandom()

def clean_str(input_str: str, check_str: str | list) -> tuple[str, list[tuple[int, chr]]]:
    letters: list[str] = []
    invalid: list[tuple[int, chr]] = []
    for index, char in enumerate(input_str):
        if char in check_str:
            letters.append(char)
        else:
            invalid.append((index, char))
    cleaned_str: str = "".join(letters)
    return cleaned_str, invalid


def insert_invalid(text_without_invalid: str, insert_chars: list[tuple[int, chr]]) -> str:
    text_list: list[str] = list(text_without_invalid)
    for position, char in insert_chars:
        text_list.insert(position, char)
    return "".join(text_list)


def encode(encode_string: str, code_str: str = "") -> tuple[str, str]:
    encode_string, invalid_chars = clean_str(encode_string, legal_chars)
    num_chars: int = len(encode_string)
    unencoded_chars: list[chr] = list(encode_string)

    if not code_str:
        rand_chars: list[chr] = [random_generator.choice(legal_chars) for _ in range(num_chars)]
        code_str = "".join(rand_chars)

    encode_nums: list[int] = [legal_chars.index(char) for char in code_str]

    encoded_str: str = ""
    for index, char in enumerate(unencoded_chars):
        num: int = legal_chars.index(char)
        num += encode_nums[index]
        num %= len(legal_chars)
        encrypted_char: chr = legal_chars[num]
        encoded_str += encrypted_char

    encoded_str = insert_invalid(encoded_str, invalid_chars)
    return encoded_str, code_str


def decode(decode_string: str, code_string: str) -> str:
    decode_string, invalid_chars = clean_str(decode_string, legal_chars)
    string_nums: list[int] = [legal_chars.index(char) for char in decode_string]
    code_nums: list[int] = [legal_chars.index(char) for char in code_string]
    decoded_str: str = ""

    for string_num, code_num in zip(string_nums, code_nums):
        decode_num: int = string_num - code_num
        character: chr = legal_chars[decode_num]
        decoded_str += character

    decoded_str = insert_invalid(decoded_str, invalid_chars)
    return decoded_str


def run(runMode: chr) -> None:
    if runMode == "e":
        chars = "".join(legal_chars)
        print(f"Only the following characters are accepted: {chars}\n")
        text: str = input("Enter the message to be encoded: ")
        code: str = "1"  # Initialize with length 1 to allow empty inputs to mean "random encryption key"
        while len(text) > len(code) > 0:
            print("Enter an encryption key, or hit Enter to get a random one. ", end="")
            code = input("It must be longer than the message: ")
        encoded_message: str
        encryption_key: str
        encoded_message, encryption_key = encode(text, code)
        print(f"Original message: {text}")
        print(f"Encrypted message: {encoded_message}")
        print(f"Encryption key: {encryption_key}")
    else:
        text = input("Enter the message to be decoded: ")
        check_text = [char for char in text if char in legal_chars]
        cypher = input("Enter the cypher to decode with: ")
        while len(cypher) < len(check_text):
            print("Cypher too short")
            cypher = input("Enter the cypher to decode with: ")
        decoded_message = decode(text, cypher)
        print(f"Encrypted message: {text}")
        print(f"Decoded message: {decoded_message}")
        print(f"Encryption key: {cypher}")


if __name__ == "__main__":
    legal_chars = [char for char in string.printable if char not in string.whitespace]
    legal_chars.append(" ")
    while True:
        mode = input("Do you want to decode or encode a message (d/e): ").lower()
        while mode not in ["d", "e"]:
            print("Enter 'e' to encode and 'd' to decode.")
            mode = input("Do you want to decode or encode message (d/e): ").lower()
        run(mode)
\$\endgroup\$
1
  • 2
    \$\begingroup\$ There are many ways in which one can use a one-time-pad in encryption. Could you expand your explanation to be more specific about your algorithm, and any other implementations it's interoperable with? Thanks. \$\endgroup\$ Commented Jan 23 at 13:33

1 Answer 1

5
\$\begingroup\$

true random

random_generator = secrets.SystemRandom()

Kudos, thank you for using the appropriate interface. (No Mersenne Twister or other PRNG.)

DRY

        mode = input("Do you ...
        while ...
            mode = input("Do you ...

It's worth drying this up a bit. Even if you change nothing else, at least extract that repeated prompt into a variable.

This wants to be a tiny helper function.

There are several ways to re-phrase the looping and the conditional. For example, you could loop while not valid_response. Then it's a matter of initializing with an invalid response. Similar remarks for your "Enter the cypher..." loop. OTOH, initializing code = "1" primed another loop very nicely.

Consider making this CLI issue someone else's problem, by using typer in conjunction with typing.Literal, or by using ArgumentParser.

chr vs str

def run(runMode: chr) -> None:

No, please don't do that.

I'm a little surprised it sneaks past mypy 1.8.0 with no errors.

It definitely fails hard at runtime when we tack on a @beartype decorator.

We're looking for type str, here.

Also, pep8 asked you nicely to spell it run_mode, and you ignored it. Probably mode would suffice.

globals

        chars = "".join(legal_chars)

Prefer to write a get_legal_chars() helper.

Or pass it in.

Or use a class and then this is self.legal_chars.

OBOB

            code = input("It must be longer than the message: ")

I don't understand that.

You meant to explain the key must be at least as long as the message, right?

Also, wouldn't you rather call the key key or otp or pad, instead of code? I would expect a code variable to maybe contain ciphertext. I guess in fairness, the trouble started with the misleading "Enter an encryption key" prompt. Names matter, they affect how we think about a problem. Choose them carefully.

implicit typing

I appreciate some of the annotations above, like text: str. We don't really need them, but it's not a bad habit to be in. Clearly mypy, and the Gentle Reader, can easily work out certain obvious types.

        encoded_message: str
        encryption_key: str

This impresses me as going slightly overboard, and as visually distracting to the flow of the source code.

Now, if we didn't have the (very nice!) -> tuple[str, str]: result annotation (in the same module!), it might be worthwhile to call out those types explicitly. Especially if mypy complained that it couldn't work out the proper type. But here? Nah.

Type annotations are tool, an aid to comprehension and one way for avoidance or early detection of bugs. But use tools where they make sense. Don't feel like there's some software component that's making you do it everywhere.

chr versus str

    unencoded_chars: list[chr] = list(encode_string)

The identifier is perfect.

But please annotate that (or rand_chars) as list[str]. This isn't C, nor Java.

Yes, you are storing individual characters in that container. But each one is a python str object. For example, let's ask about the SPACE character.

>>> type(chr(32))
<class 'str'>

Yup, that character is definitely a string.

standard terminology

Cryptographers have standard terms to describe crypto concepts.

Instead of encode_str + encoded_str, prefer cleartext + ciphertext.

time complexity

        num: int = legal_chars.index(char)

This is a moderately expensive lookup. You're asking the cPython interpreter to repeatedly do a linear scan of that string. Now, it is C code doing the looping, not interpreted bytecode, but still. Consider building a dict beforehand, so within the loop we'll see just O(1) constant time lookups.

Similar remarks apply when your're building encode_nums.

Anyway, kudos, the num += and num %= expressions are extremely clear, and I thank you for them. (I still have to mentally translate that encode_nums is really "the pad".)

Similar remarks when decoding:

    string_nums: list[int] = [legal_chars.index(char) for char in decode_string]
    code_nums:   list[int] = [legal_chars.index(char) for char in code_string]

There's many ways to do this efficiently. As an alternative to dict, you might find str.maketrans() to be of interest.

simple identifiers

def insert_invalid(text_without_invalid: str, insert_chars: list[tuple[int, chr]]) -> ...

Wow, that first parameter takes a moment to wrap your head around. Why? Because it ain't got no single negative one time, no no it's double negative.

Humans have a tougher time with negated identifiers than with positive ones. Prefer valid_text here.

And zomg is that final type annotation ever helpful, thank you so much. Perfect!

The insertion algorithm is nice enough, it looks correct. But perhaps it, and the list of tuples, are fancier than needed? If you were to tackle this encryption task again, you might prefer to just leave (unencoded) bad characters in the middle of ciphertext as you iterate along.

I do really appreciate that you broke this out as a separate helper, very nice.

Hmmm, interesting, this helper has two callers, for the two directions. Cool.

alignment

        print(f"Original message: {text}")
        print(f"Encrypted message: {encoded_message}")
        print(f"Encryption key: {encryption_key}")

nit: If you insert just a handful of SPACE characters after : colons, then it should be visually apparent to the user that cleartext, ciphertext, and OTP all have the identical length. (Or it will be apparent how many characters from end of pad went unused.)

tricky code

        character: chr = legal_chars[decode_num]

The detail here is we might be chasing a negative index, and then the Right Thing will still happen.

Consider calling this out in a # comment.

Consider using explicit % modulo, for parallel construction with the encryption operation, so we have non-negative indexes.

If someone has to maintain this code a few months down the road, or has to translate it to a non-python language, there's no need to leave them scratching their head about how this ever managed to work in the first place.

automated tests

"This is an example" encoded with "Mary had a little lamb" gives "8rJ!~zCc9x~zPDPKzd".

Thank you for that.

It would have been even nicer if wrapped up within a test suite.

summary

You clearly went to a lot of trouble to come up with well-considered identifiers, and to break out single responsibility helpers where needed. That is all for the good, keep it up! And in some cases those optional type annotations proved super helpful.

I would say, outside of the signatures, don't go overboard with redundant type annotations. Often the derived type is pretty obvious, both to man and machine. Where it isn't, then by all means throw an extra annotation in there.

You have a tendency to accept an iterable or a container (like str), and explode it out into a list. I predict that you soon will be comfortable enough with such loops and transformations that you will no longer feel a need to always name and create a new container for an existing container. Put another way: if there's a name for a concept in the Business Domain, then by all means introduce a suitably named temp var for it, but if not then maybe an unnamed intermediate expression suffices.

practices for handling keying material

The well-known downside of OTP is that, while it lets you send a giant message, now you have a giant pad to carry around and securely disclose to the recipient. And reusing the pad Would Be Bad.

In the code and the review context I saw no warnings or advice about how to properly handle sensitive pad material.

Typically we would generate "a lot" of pad, maybe hand it off in person, on a thumb drive, to the intended recipient before they get on a plane, and then send several "short" messages. We would take care not to exhaust the pad. Yes, the current code could support such a scenario, but it would be clumsy. We don't store the pad in a file, and there's currently no support for seeking to appropriate offset for the "next" portion of the pad that we should use. There's no support for storing, and then burning, portions of the pad, to help prevent accidental reuse.

A pure CLI approach might be more convenient, and more composable, than the current interactive prompting.

So there's room for feature requests, if you want to keep playing with this.


This codebase achieves its design goals.

I would be willing to delegate or accept maintenance tasks on it.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.