Encodes/decodes a message using the Rail Fence Cipher

rails.py

# Author      : Lokno Ketchup
# Description : Encodes/decodes a message using the Rail Fence Cipher
#     Allows for recursively encoding of a message multiple times.
# usage: rails.py [-h] [-d] [-b BLOCKS] input output count [count ...]
# Parameters  :
#     [-d,--decode]     - decodes if present (defaults to encoding)
#     [-b BLOCKS, --blocks BLOCKS] - expects encoded message in groupings of the specified size
#     <input>           - file to encode/decode
#     <output>          - file to write output.
#     count [count ...] - list of rail counts to encode
#          or decode a message. List in order to encode the message
#          from left to right, even when requesting decoding.

import sys
import re
import argparse

# Name: rfc_encode
# Description: Encodes a string s using the Rail
#    Fence Cipher using railCount rails.
#
# Parameters:
# s         - string to encode
# railCount - number of rails to encode across
# blockSize - IF nonzero, the input string s
#     is stripped of nonalphanumberic characters
#     and the encoded message is printed in space
#     separated blocks of size blockSize;
#     OTHERWISE (blockSize=0) s is encoded as is
def rfc_encode( s, railCount, blockSize=0 ):

    if railCount < 2:
        railCount = 2

    if blockSize != 0:
        # remove all non-alphanumeric characters
        s = s.upper()
        s = re.sub("[^A-Z0-9]",'',s)
        
    l = len(s)
    rails = [''] * railCount
    count = 0

    mo = railCount*2-2
    ro = railCount-1
    for c in s:
        idx = ro-abs((count % mo)-ro)
        rails[idx] += c
        count += 1
        
    message = ''.join(rails)

    if blockSize != 0:
        blocks = [message[i:i+blockSize] for i in range(0, l, blockSize)]
        message = ' '.join(blocks)
        
    return message

# Name: rfc_decode
# Description: Decodes a string s using the Rail
#    Fence Cipher using railCount rails.
#
# Parameters:
# s         - string to decode
# railCount - number of rails to decode across
# inBlocks - IF true, string s is assumed to be
#     divided into space separated blocks.
def rfc_decode( s, railCount, inBlocks=False ):
    
    if railCount < 2:
        railCount = 2
        
    if inBlocks:
        s = s.replace(' ','')
    
    l = len(s)
    message = [' '] * l

    mo = railCount*2-2
    ro = railCount-1

    sidx = 0      # moving start index in output 
    cidx = 0      # current index in output
    rail = 0      # current rail 
    step = [0,0]  # offset sequence for advancing through message length
    stepIdx = 0   # current index in step
    
    cidx = sidx
    step[0] = (mo-(rail % ro)*2)
    step[1] = mo-step[0]
    if step[1] == 0:
        step[1] = step[0]
    
    for c in s:
        message[cidx] = c
        cidx += step[stepIdx]
        stepIdx = 1-stepIdx
        if cidx >= l:
            sidx += 1
            rail += 1
            stepIdx = 0
            cidx = sidx
            step[0] = (mo-(rail % ro)*2)
            step[1] = mo-step[0]
            if step[1] == 0:
                step[1] = step[0]
    return ''.join(message)

parser = argparse.ArgumentParser(description='''Encodes/decodes a message using the Rail Fence Cipher. 
                    Allows for recursively encoding of a message multiple times.''')   
parser.add_argument('-d','--decode', action='store_true', help='decodes message if flag is present')
parser.add_argument('-b','--blocks', type=int, default=0, help='expects encoded message in groupings of the specified size')
parser.add_argument('input', help='file to encode/decode' )
parser.add_argument('output', help='output file' ) 
parser.add_argument('railCounts', metavar='count', type=int, nargs='+', help='''rail count(s) to encode 
                    or decode a message. List multiple counts in the order in which to encode the message 
                    from left to right, even when requesting decoding.''') 

args = parser.parse_args()
with open(args.input,'rb') as fin:
    with open(args.output,'wb') as fout:
        s = fin.read()
        if args.decode:
            args.railCounts = args.railCounts[::-1]
            for count in args.railCounts:
                s = rfc_decode(s, int(count), args.blocks)
        else:           
            for count in args.railCounts:
                s = rfc_encode(s, int(count), args.blocks)
        fout.write(s)