polybius.py

#!/usr/bin/env python
import sys
import optparse

"""polybius.py: Encrypt and decrypt using the polybius cipher with a 6 by 6 key matrix."""

__author__ = "Christian Engel"
__license__ = "MIT"
__version__ = "1.0.0"
__maintainer__ = "Christian Engel"


class Polybius:
    """Implements Polybius cipher."""

    def __init__(self):
        self.enc_map = {}
        self.dec_map = {}

    def setKey(self, key: str) -> bool:
        """Initializes both key dictionaries and returns True if successful.
        key: Key string. Must have 36 characters.

        Note: If the key is invalid, will return False.
        """
        if len(key) != 36:
            return False

        i = 1
        j = 1

        for k in key:
            if k in self.enc_map:
                return False
            c = str(i) + str(j)
            self.enc_map[k] = c
            self.dec_map[c] = k

            if j == 6:
                j = 1
                i += 1
            else:
                j += 1
        return True

    def encrypt(self, plain: str) -> str:
        """Returns a sequence of two-digit numbers by encrypting a plain text.

        plain: The plaintext passed as single string.
        """
        plain = plain.lower()
        cipher = ""
        for c in plain:
            if c in self.enc_map:
                cipher += self.enc_map[c] + " "
        return cipher

    def decrypt(self, cipher: str) -> str:
        """Returns a decrypted text based on the input sequence.

        cipher: The sequence of two-digit numbers.
        """
        plain = ""
        for i in range(len(cipher) - 1):
            c = cipher[i] + cipher[i + 1]
            if c in self.dec_map:
                plain += self.dec_map[c]
        return plain

    def print(self):
        """Prints the encryption map."""
        print("  | 1 2 3 4 5 6")
        print("--+---------------")
        i = 1
        for i in range(1, 7):
            for j in range(1, 7):
                if j == 1:
                    print(str(i) + " | ", end="")
                c = str(i) + str(j)
                print(self.dec_map[c] + " ", end="")
                if j == 6:
                    print("\n", end="")


if __name__ == "__main__":
    parser = optparse.OptionParser(
        '\n\t./polybius.py -e "HELLO WORLD" -k "abcdefghijklmnopqrstuvwxyz0123456789"\n\t./polybius.py -d "KHOORZRUOG" -k "abcdefghijklmnopqrstuvwxyz0123456789"'
    )
    parser.add_option(
        "-e", "--encrypt", dest="plain", type="string", help="encrypt mode"
    )
    parser.add_option(
        "-d", "--decrypt", dest="cipher", type="string", help="decrypt mode"
    )
    parser.add_option("-k", dest="k", type="string", help="secret key")
    options, args = parser.parse_args()

    if options.k == None or (options.cipher == None and options.plain == None):
        print(parser.usage)
        sys.exit(1)

    k = options.k
    plain = options.plain
    cipher = options.cipher

    pb = Polybius()
    success = pb.setKey(k)

    if not success:
        print("INVALID KEY")
        print(parser.usage)
    elif plain != None:
        cipher = pb.encrypt(plain.lower())
        print(cipher.upper())
    else:
        plain = pb.decrypt(cipher.lower())
        print(plain.upper())