vigenere.py

from __future__ import unicode_literals


def vigenere_encode(m, k):
    """
    >>> vigenere_encode('A', 'A')
    'A'
    >>> vigenere_encode('SUPER SECRET', 'KEY')
    'CYNOV QOGPOX'
    >>> vigenere_encode('LOREM IPSUM DOLOR SIT AMET CONSECTETUR ADIPISCING ELIT', 'SECRET')
    'DSTVQ BHWWD HHDST JMM SQGK GHFWGTXXLYT RHBHMUTMGY INZX'

    :type m: str
    :type k: str
    :rtype: str
    """
    m = m.upper()
    k = k.upper()
    key_position = 0
    c = ''

    for m_char in m:
        m_index = ord(m_char) - 65
        key_length = len(k)

        if not 0 <= m_index < 65:
            c_char = m_char

        else:
            k_char = k[key_position]
            k_index = ord(k_char) - 65
            c_index = (m_index + k_index) % 26
            c_char = chr(c_index + 65)
            key_position = (key_position + 1) % key_length

        c += c_char

    return c


def vigenere_decode(c, k):
    """
    >>> vigenere_decode('A', 'A')
    'A'
    >>> vigenere_decode('CYNOV QOGPOX', 'KEY')
    'SUPER SECRET'
    >>> vigenere_decode('DSTVQ BHWWD HHDST JMM SQGK GHFWGTXXLYT RHBHMUTMGY INZX', 'SECRET')
    'LOREM IPSUM DOLOR SIT AMET CONSECTETUR ADIPISCING ELIT'

    :type m: str
    :type k: str
    :rtype: str
    """
    c = c.upper()
    k = k.upper()
    key_position = 0
    m = ''

    for c_char in c:
        c_index = ord(c_char) - 65
        key_length = len(k)

        if not 0 <= c_index < 65:
            m_char = c_char

        else:
            k_char = k[key_position]
            k_index = ord(k_char) - 65
            m_index = (c_index - k_index) % 26
            m_char = chr(m_index + 65)
            key_position = (key_position + 1) % key_length

        m += m_char

    return m