a bit of a clean-up

1 files changed, 123 insertions, 56 deletions

diff --git a/ls47.py b/ls47.py
index 8736230..ae8b201 100755
--- a/ls47.py
+++ b/ls47.py
@@ -6,32 +6,35 @@
 # Originally written by Mirek Kratochvil (2017)
 # Python3 port by Bernhard Esslinger (Feb 2018)
 
+import argparse
 import random
+import sys
 
-letters = "_abcdefghijklmnopqrstuvwxyz.,:?!'()0123456789+-*/"
+letters = "_ABCDEFGHIJKLMNOPQRSTUVWXYZ.,:?!'()0123456789+-*/"
 tiles = list(zip(letters, map(lambda x: (x // 7, x % 7), range(7 * 7))))
-padding_size = 10
+padding_size_min = 5
+padding_size_max = 15
 
 
 def check_key(key):
     if len(key) != len(letters):
-        raise ValueError('Wrong key size')
+        raise ValueError("Wrong key size")
     cnts = {}
     for c in letters:
         cnts[c] = 0
     for c in key:
         if not c in cnts:
-            raise ValueError('Letter ' + c + ' not in LS47!')
+            raise ValueError(f"Letter {c} not in LS47!")
         cnts[c] += 1
         if cnts[c] > 1:
-            raise ValueError('Letter ' + c + ' duplicated in key!')
+            raise ValueError(f"Letter {c} duplicated in key!")
 
 
 def find_ix(letter):
     m = [l for l in tiles if l[0] == letter]
     if len(m) != 1:
-        raise ValueError('Letter ' + letter + ' not in LS47!')
-    for (l, pos) in m:
+        raise ValueError(f"Letter {letter} not in LS47!")
+    for l, pos in m:
         return pos
 
 
@@ -39,7 +42,7 @@ def find_pos(key, letter):
     p = key.find(letter)
     if p >= 0 and p < 7 * 7:
         return (p // 7, p % 7)
-    raise ValueError('Letter ' + letter + ' not in key?!')
+    raise ValueError(f"Letter {letter} not in key!")
 
 
 def add_pos(a, b):
@@ -55,19 +58,19 @@ def find_at_pos(key, coord):
 
 
 def rotate_right(key, row, n):
-    mid = key[7 * row:7 * (row + 1)]
+    mid = key[7 * row : 7 * (row + 1)]
     n = (7 - n % 7) % 7
-    return key[:7 * row] + mid[n:] + mid[:n] + key[7 * (row + 1):]
+    return key[: 7 * row] + mid[n:] + mid[:n] + key[7 * (row + 1) :]
 
 
 def rotate_down(key, col, n):
-    lines  = [key[i * 7:(i + 1) * 7] for i in range(7)]
-    lefts  = [l[:col] for l in lines]
-    mids   = [l[col] for l in lines]
-    rights = [l[col + 1:] for l in lines]
+    lines = [key[i * 7 : (i + 1) * 7] for i in range(7)]
+    lefts = [l[:col] for l in lines]
+    mids = [l[col] for l in lines]
+    rights = [l[col + 1 :] for l in lines]
     n = (7 - n % 7) % 7
     mids = mids[n:] + mids[:n]
-    return ''.join(lefts[i] + mids[i] + rights[i] for i in range(7))
+    return "".join(lefts[i] + mids[i] + rights[i] for i in range(7))
 
 
 def rotate_marker_right(m, row, n):
@@ -85,25 +88,21 @@ def rotate_marker_down(m, col, n):
 
 
 def derive_key(password):
-    k = letters
-    for c in password:
-        (crow, ccol) = find_pos(k, c)
-        (rows1, cols1) = find_ix(c)
-        k = rotate_right(k, crow, cols1)
-        (crow, ccol) = add_pos((crow, ccol), (0, cols1))
-        k = rotate_down(k, ccol, rows1)
-        (crow, ccol) = add_pos((crow, ccol), (rows1+1, 1))
-        (rows2, cols2) = find_ix(find_at_pos(k, (crow, ccol)))
-        k = rotate_right(k, crow, cols2)
-        (crow, ccol) = add_pos((crow,ccol), (0, cols2))
-        k = rotate_down(k, ccol, rows2)
-    return k
+    k = ''.join(reversed(letters))
+    i = 0
+    for c in password + ' ':
+        ix = k.find(c,i)
+        if ix<0:
+            continue
+        k = k[:i] + k[ix] + k[i:ix] + k[ix+1:]
+        i += 1
+    return ''.join(reversed(k))
 
 
 def encrypt(key, plaintext):
     check_key(key)
     mp = (0, 0)
-    ciphertext = ''
+    ciphertext = ""
     for p in plaintext:
         pp = find_pos(key, p)
         mix = find_ix(find_at_pos(key, mp))
@@ -121,7 +120,7 @@ def encrypt(key, plaintext):
 def decrypt(key, ciphertext):
     check_key(key)
     mp = (0, 0)
-    plaintext = ''
+    plaintext = ""
     for c in ciphertext:
         cp = find_pos(key, c)
         mix = find_ix(find_at_pos(key, mp))
@@ -137,34 +136,102 @@ def decrypt(key, ciphertext):
 
 
 def encrypt_pad(key, plaintext, signature):
-
-    # TODO it would also be great to randomize the message length.
-
     check_key(key)
-    padding = ''.join(map(lambda x: letters[random.randint(0,
-                      len(letters) - 1)], range(padding_size)))
+    padding = "".join(
+        map(
+            lambda x: letters[random.randint(1, len(letters) - 1)],
+            range(random.randint(padding_size_min, padding_size_max)),
+        )
+    )
 
-    return encrypt(key, padding + plaintext + '---' + signature)
+    return encrypt(key, padding + "_" + plaintext + "_" + signature)
 
 
-def decrypt_pad(key, ciphertext):
+def decrypt_pad(key, ciphertext, signature):
     check_key(key)
-    return decrypt(key, ciphertext)[padding_size:]
-
-
-if __name__ == '__main__':
-
-    # a bit of test!
-
-    print('letters in this implementation (line by line):')
-    print(letters)
-    # print('tiles positions: ' + str(tiles))
-
-    key = derive_key('a')
-    print('test key: ' + key)
-    enc = encrypt_pad(key, 'conflagrate_the_rose_bush_at_six!',
-                      'peace-vector-3')
-
-    print('encrypted test: ' + enc)
-    dec = decrypt_pad(key, enc)
-    print('decrypted test: ' + dec)
+    padded = decrypt(key, ciphertext)
+    pad_end = padded.find("_")
+    if pad_end < padding_size_min or pad_end > padding_size_max:
+        raise ValueError("mangled padding")
+    padded = padded[pad_end + 1 :]
+    if len(signature) == 0:
+        return padded
+    if padded[-len(signature) - 1 :] != "_" + signature:
+        raise ValueError("bad signature")
+    return padded[: -len(signature) - 1]
+
+
+if __name__ == "__main__":
+
+    ap = argparse.ArgumentParser(
+        prog=sys.argv[0],
+        description="Computer implementation of LS47.",
+        epilog="This is flimsy! Encrypt physically!",
+    )
+
+    ap.add_argument(
+        "-d", "--decrypt", dest="decrypt", action="store_true", help="run decryption"
+    )
+    ap.add_argument(
+        "-p",
+        "--pad-and-sign",
+        dest="sign",
+        help="On encryption, add a random padding and a signature at the end. On decryption, expect a random padding and the exact signature. Specify empty string to attach no (or accept any) signature.",
+    )
+    ap.add_argument(
+        "-s",
+        "--password",
+        dest="password",
+        help="Password string (to be expanded by the simple swapping algorithm)",
+    )
+    ap.add_argument(
+        "-K",
+        "--full-key",
+        dest="fullkey",
+        help="Shared key (the initial state of the board, line by line, always with `_' instead of space; used only if password is not specified)",
+    )
+    ap.add_argument(
+        "-I",
+        "--exact-input",
+        dest="inx",
+        help="Do not convert input spaces to `_', ignore them instead.",
+        action="store_true"
+    )
+    ap.add_argument(
+        "-O",
+        "--exact-output",
+        dest="outx",
+        help="Do not convert output `_'s to spaces.",
+        action="store_true"
+    )
+
+    args = ap.parse_args()
+
+    if args.password:
+        key = derive_key(args.password if args.inx else args.password.replace(' ', '_'))
+    elif args.fullkey:
+        key = args.fullkey
+    else:
+        raise ValueError("Specify at least one of password and full key (options -s and -K)")
+
+    input = sys.stdin.read()
+    if not args.inx:
+        input = input.replace(' ', '_')
+    letter_set = set(letters)
+    input = ''.join(filter(lambda x:x in letter_set, input.upper()))
+
+    if args.decrypt:
+        if args.sign is None:
+            output = decrypt(key, input)
+        else:
+            output = decrypt_pad(key, input, args.sign)
+    else:  # encrypt
+        if args.sign is None:
+            output = encrypt(key, input)
+        else:
+            output = encrypt_pad(key, input, args.sign)
+
+    if not args.outx:
+        output = output.replace('_', ' ')
+
+    print(output)