Little rail fence cipher

rail_fence_cipher.py

#!/usr/bin/env python
# Minified rail fence cipher (https://en.wikipedia.org/wiki/Rail_fence_cipher)
# Matt Borgerson, 2018
import argparse

# Iterative approach where we step through the input and add
# each character to the corresponding rail, then combine the
# rails in the end.
def encode(msg, n):
    rm,l = list(range(n))+list(range(n-2,0,-1)),['']*n
    for i in range(len(msg)): l[rm[i%len(rm)]] += msg[i]
    return ''.join(l)

# Iterative approach where we first extract the rail text, then
# successively pop from each rail to descramble the text.
def decode(msg, n):
    rm,l,r,s = list(range(n))+list(range(n-2,0,-1)),[0]*n,['']*n,0
    for i in range(len(msg)): l[rm[i%len(rm)]] += 1
    for i in range(n): r[i],s=list(msg[s:s+l[i]]),s+l[i]
    return ''.join([r[rm[i%len(rm)]].pop(0) for i in range(len(msg))])

msg, n = 'WEAREDISCOVEREDFLEEATONCE', 3
assert(decode(encode(msg, n), n) == msg)

def main():
    ap = argparse.ArgumentParser('Rail Fence Encoder/Decoder')
    ap.add_argument('msg')
    ap.add_argument('--num_rails', '-n', type=int, required=False)
    ap.add_argument('--decode', '-d', action='store_true')
    ap.set_defaults(num_rails=3)
    args = ap.parse_args()
    func = decode if args.decode else encode
    print(func(args.msg, args.num_rails))

if __name__ == '__main__':
    main()