#!/usr/bin/env python3
# InterPlay-36: reference Python 3 implementation for the
# Interleaved 6x6 Playfair crypto algorithm
#
# Features:
# * fixed 6x6 alphabet (whitespace, 1-9, a-z; 0 is replaced with o)
# * normal Playfair rules for table keying and encryption/decryption
# * the plaintext is interleaved with random chars
#   before encryption and they are removed after decryption
# * since the cryptogram may begin with a space, a random prefix is prepended
# * (its length depends on the key length)
#
# Created by Luxferre in 2025, released into public domain

import sys, math, secrets
from collections import OrderedDict

pfgrid = ' 123456789abcdefghijklmnopqrstuvwxyz'
pfsize = int(math.sqrt(len(pfgrid)))

# message/keystring preparation step: replace 0 with o then filter out all invalid chars
def filtermsg(msg:str):
    return ''.join(list(filter(lambda i: i in pfgrid, msg.lower().replace('0', 'o').replace('\n', '   '))))

# grid keying algorithm
def keygrid(key:str):
    # start with the standard fill-in
    startgrid = list(OrderedDict.fromkeys(filtermsg(key) + pfgrid).keys())
    # extract the transposition keys
    tkey1 = []
    tkey2 = []
    rkey = list(map(lambda x: pfgrid.index(x), startgrid[0:pfsize])) # first row
    srkey = sorted(rkey)
    for c in rkey:
        tkey1.append(srkey.index(c))
    rkey = list(map(lambda x: pfgrid.index(x), startgrid[pfsize:pfsize*2])) # second row
    srkey = sorted(rkey)
    for c in rkey:
        tkey2.append(srkey.index(c))
    # perform the zigzag row transposition into the intermediate grid
    igrid = []
    for row in tkey1:
        for i in range(0, pfsize):
            igrid.append(startgrid[((row + (i&1)) % pfsize) * pfsize + i])
    # perform the straight columnar transposition into the final key grid
    keygrid = []
    for col in tkey2:
        for i in range(0, pfsize):
            keygrid.append(igrid[i * pfsize + col])
    return keygrid

# get character coordinates in the grid
def getcoords(grid:list, char:str):
    pos = grid.index(char) 
    return (pos % pfsize, int(pos // pfsize))

# get grid character by coordinates
def getgridchar(grid:list, x:int, y:int):
    return grid[y * pfsize + x]

# Interleaved Playfair encryption method
def ipfencrypt(msg:str, key:str):
    kg = keygrid(key)
    # prepare the message
    msg = filtermsg(msg).strip()
    # perform the encryption
    enc = []
    for c in msg: 
        # shape the digraphs by interleaving random characters
        nc = c
        while nc == c: # the second char in digraph must differ
            nc = secrets.choice(pfgrid)
        # run the algo
        x1, y1 = getcoords(kg, c)
        x2, y2 = getcoords(kg, nc)
        if x1 != x2 and y1 != y2: # no coordinates are equal
            enc.append(getgridchar(kg, x2, y1))
            enc.append(getgridchar(kg, x1, y2))
        else:
            if x1 == x2: # same column: get the coords below
                y1 = (y1 + 1) % pfsize
                y2 = (y2 + 1) % pfsize
            elif y1 == y2: # same row: get the coords to the right
                x1 = (x1 + 1) % pfsize
                x2 = (x2 + 1) % pfsize
            # swap the ciphertext character
            enc.append(getgridchar(kg, x2, y2))
            enc.append(getgridchar(kg, x1, y1))
    # generate a random prefix
    preflen = (len(key) % 10) + 1
    prefix = ''
    for i in range(preflen):
        prefix += secrets.choice(pfgrid[1:])
    return prefix + ''.join(enc)

# Interleaved Playfair decryption method
def ipfdecrypt(enc:str, key:str):
    kg = keygrid(key)
    # remove the prefix, then add a space if the cryptogram
    # had been stripped and ended with a space
    preflen = (len(key) % 10) + 1
    enc = enc[preflen:].rstrip()
    if len(enc) & 1 == 1:
        enc += ' '
    # split into digraphs
    digraphs = [enc[i:i+2] for i in range(0, len(enc), 2)]
    # perform the decryption
    msg = []
    for dg in digraphs: # for IPF, we only decrypt the first digraph char
        x1, y1 = getcoords(kg, dg[0])
        x2, y2 = getcoords(kg, dg[1])
        if x1 != x2 and y1 != y2: # no coordinates are equal
            msg.append(getgridchar(kg, x2, y1))
        else:
            if x1 == x2: # same column: get the coords above
                y2 -= 1
                if y2 < 0:
                    y2 += pfsize
            elif y1 == y2: # same row: get the coords to the left
                x2 -= 1
                if x2 < 0:
                    x2 += pfsize
            msg.append(getgridchar(kg, x2, y2))
    # finalize the decrypted message
    return ''.join(msg).replace('   ', '\n').strip()

# entrypoint parameters: [mode] [keystring]
# modes: e - encrypt with 6x6 IPF, d - decrypt with 6x6 IPF
def main():
    mode = sys.argv[1].lower()
    key = sys.argv[2]
    msg = sys.stdin.read()
    if mode.startswith('e'):
        print(ipfencrypt(msg, key))
    else:
        print(ipfdecrypt(msg, key))

if __name__ == '__main__':
    main()