otp.py

import secrets
import string
from typing import Sequence

"""
This is a set of functions for demonstrating one-time pad
encryption using a twenty-six letter alphabet.

This code is a companion to the article at
https://notes.ericjiang.com/posts/1051
"""

def add_text(text1: str, text2: str) -> str:
    """
    >>> add_text('HELLO', 'ABCDE')
    'HFNOS'
    """
    return "".join([
        number_to_char(
            (char_to_number(c1) + char_to_number(c2)) % 26
        )
        for c1, c2 in zip(text1, text2)])


def subtract_text(text1: str, text2: str) -> str:
    """
    Given two strings of equal length, subtract each letter in `text2`
    from the corresponding letter in `text1`.

    >>> subtract_text('HELLO', 'ABCDE')
    'HDJIK'
    """
    return "".join([
        number_to_char(
            (char_to_number(c1) - char_to_number(c2)) % 26
        )
        for c1, c2 in zip(text1, text2)])


def strip_spaces(text: str) -> str:
    """
    Removes all space characters from a string.

    >>> strip_spaces("Hello world")
    'Helloworld'
    """
    return text.replace(" ", "")


def random_key(length: int) -> str:
    """
    Creates a random string of letters of `length` characters.

    >>> random_key(5)
    'CKMML'
    """
    return ''.join(secrets.choice(string.ascii_uppercase) for _ in range(length))


def char_to_number(char: str) -> int:
    char = char.upper()

    if ord(char) < ord('A') or ord(char) > ord('Z'):
        raise ValueError("character out of range")

    return ord(char) - ord('A')


def number_to_char(num: int) -> str:
    if num < 0 or num > 25:
        raise ValueError("number out of range")

    return chr(num + ord('A'))


def text_to_nums(text: str) -> Sequence[int]:
    return [char_to_number(x) for x in text]


def nums_to_text(nums: Sequence[int]) -> str:
    return "".join([number_to_char(x) for x in nums])