#@# Classical cryptography
#@# ======================


#@# ---------------------------------------------------------------------------
#@# Exercise 1: shift cipher
#@# [with solution]
#@#
#@# Implement the shift cipher: shift each letter of the message (m) by a fixed
#@# amount (= key k), wrapping around at the end of the alphabet. For example,
#@# the Caesar cipher is a shift cipher with k = 3, i.e.:
#@#    A -> D, B -> E, C -> F, ..., W -> Z, X -> A, Y -> B, Z -> C,
#@# so it encrypts m="VERYSECRET" as c="YHUBVHFUHW".
#@# The message (m) contains only uppercase English letters (A, B, ..., Z),
#@# and the secret key (k) is an integer between 1 and 25, the shift amount.

# Convert between letters and their indices in the alphabet:
def let2ind (letter): # "A" -> 0, "B" -> 1, ..., "Z" -> 25
    return ord(letter) - 65
def ind2let (index):  # 0 -> "A", 1 -> "B", ..., 25 -> "Z"
    return chr(index + 65)

# Encryption
def shift_enc (m, k):
    c = ""
    for i in range(len(m)):
        x = let2ind(m[i])
        x = (x + k) % 26  # '% 26' does the "wrapping around"
        c += ind2let(x)
    return c

# Decryption
def shift_dec (c, k):
    m = ""
    for i in range(len(c)):
        x = let2ind(c[i])
        x = (x - k) % 26
        m += ind2let(x)
    return m

# Generate a key
def shift_gen ():
    # Generate a random number between 1 and 25
    # (omiting k = 0, which wouldn't do anything)
    return randrange(1, 26)  # Sage-only function

#@# Testing
#@# (Note: "#@#" means "don't change this line!".)
def shift_test (): #@#
    print("1. Shift cipher") #@#
    m, c = "VERYSECRET", "YHUBVHFUHW" #@#
    test_cipher(shift_enc, shift_dec, 3, m, c) #@#
    k = shift_gen() #@#
    m = "THISISAVERYSECRETMESSAGE" #@#
    test_cipher(shift_enc, shift_dec, k, m) #@#
    print() #@#

#@# Test any cipher (for all exercises)
def test_cipher (enc, dec, k, m, c=None): #@#
    c_ = enc(m, k) #@#
    error = (c and c_ != c) or c_ == None #@#
    print(m, "->", c_, end="") #@#
    if error: print(" (ERROR!)", end="") #@#
    print() #@#
    if c_ == None: return #@#
    if error and c: #@#
        print(" "*len(m)+"(-> "+c+")") #@#
    m_ = dec(c_, k) #@#
    if m != m_: #@#
        print(m_, "<- (decryption ERROR!)") #@#

#shift_test()

#@# Bonus question:
#@# How to fill '...' below to make it equivalent to shift_dec() above?

#def shift_dec (c, k):
    #return shift_enc(c, ...)


#@# ---------------------------------------------------------------------------
#@# Exercise 2: brute-force attack
#@#
#@# The ciphertext (c) given below was encrypted by the shift cipher with some
#@# unknown key k. Break it using brute-force attack, i.e. try to decrypt it
#@# with all possible keys, and (manually) select the plaintext that makes
#@# sense as an English text.
c = "SEDWHQJKBQJYEDIOEKXQLUTUSYFXUHUTJXYICUIIQWU" #@#

def shift_brute_force (c):
    pass #TODO (<- replace it with your code)

## Testing (uncomment when ready):
#print("2. Brute-force attack")
#shift_brute_force(c)
#print()

#@# Note 1: It can be done in a more automated way, see the next class.
#@# Note 2: To resist brute-force attack, we need a bigger key space.


#@# ---------------------------------------------------------------------------
#@# Exercise 3: mono-alphabetic substitution cipher
#@#
#@# Instead of shifting each letter by the same amount, the key is a random
#@# permutation that defines which letter will be replaced by which letter.
#@# For example:
#@#    A B C D E F G H I J K L M N O P Q R S T U V W X Y Z  (plaintext letters)
#@#    X E U A D N B K V M R O C Q F S Y H W G L Z I J P T  (ciphertext letters)
#@# We encode this permutation by storing the letters in the second line in a
#@# list, or rather their alphabet index between 0 and 25 (let2ind()):
#@#    k = [23,4,20,0,3,13,1,10,21,12,17,14,2,16,5,18,24,7,22,6,11,25,8,9,15,19]
#@# Implement encryption and decryption for this cipher.

# Encryption
def mono_enc (m, k):
    pass #TODO (your code here)

# Decryption
def mono_dec (c, k):
    pass #TODO (your code here)

# Generate a key (random permutation)
def mono_gen ():
    k = [i for i in range(26)]  # [0, 1, ..., 25]
    shuffle(k)
    return k

#@# Testing
def mono_test (): #@#
    print("3. Mono-alphabetic substitution") #@#
    k = mono_gen() #@#
    m = "SECRETMESSAGE" #@#
    test_cipher(mono_enc, mono_dec, k, m) #@#
    #@# Shift cipher is just a special case:
    k = [(i + 3) % 26 for i in range(26)] #@#
    m, c = "VERYSECRET", "YHUBVHFUHW" #@#
    test_cipher(mono_enc, mono_dec, k, m, c) #@#
    print() #@#
#mono_test()  # (<- uncomment when ready)

#@# Bonus questions:
#@# 1. How big is the key space (i.e. how many possible keys are there)?
#@# 2. How long would it take to brute-force this cipher?


#@# ---------------------------------------------------------------------------
#@# Exercise 4: Vigenère cipher (or poly-alphabetic shift cipher)
#@#
#@# Another way to improve the shift cipher is to use different shifts in each
#@# position. The key is a string that is repeated to fill the plaintext, and
#@# each letter is shifted by the amount defined by the key's current letter.
#@# For example, for k = "DECK", the first letter is shifted by +3 (for D),
#@# the second one by +4 (for E), and so on:
#@#     m:  SECRETMESSAGE
#@#     k*: DECKDECKDECKD
#@#     c:  VIEBHXOOVWCQH
#@# Implement encryption and decryption for the Vigenère cipher.

# Encryption
def vigenere_enc (m, k):
    pass #TODO

# Decryption
def vigenere_dec (c, k):
    pass #TODO

#@# Testing
def vigenere_test (): #@#
    print("4. Vigenère cipher") #@#
    k = "DECK" #@#
    m, c = "SECRETMESSAGE", "VIEBHXOOVWCQH" #@#
    test_cipher(vigenere_enc, vigenere_dec, k, m, c) #@#
    #@# Shift cipher is just a special case:
    m, c = "VERYSECRET", "YHUBVHFUHW" #@#
    k = "D" #@# for k = 3
    test_cipher(vigenere_enc, vigenere_dec, k, m, c) #@#
    print() #@#
#vigenere_test()


#@# ---------------------------------------------------------------------------
#@# Exercise 5: Grille cipher
#@#
#@# The grille cipher works by writing the plaintext message on a grid through
#@# a square paper with random holes in it, and then filling the remaining
#@# grid positions by random letters. For example, encrypting "MESSAGE":
#@#       # . # # .       # M # # E       V M W K E
#@#       # # . # #       # # S # #       X H S Q U
#@#       . # # # .  -->  S # # # A  -->  S Y M N A
#@#       # # # . #       # # # G #       N D W G P
#@#       # . # # #       # E # # #       S E N L N
#@# The resulting grid of letters (the ciphertext) can be decrypted by
#@# placing the exact same grille back to the grid, and reading the letters.
#@# More info: https://en.wikipedia.org/wiki/Grille_(cryptography)
#@# Implement the grille cipher. The key (the paper with holes) is encoded as
#@#    k = (n, [(x1,y1), (x2,y2), ...]),
#@# where n is the size of the square (in this case, n = 5), and (x1,y1) etc.
#@# are the positions of the holes (the top left being (0,0)).
#@# So in the above example:
#@#    k = (5, [(1,0), (4,0), (2,1), (0,2), (4,2), (3,3), (1,4)])
#@#    m = "MESSAGE"
#@#    c = "VMWKEXHSQUSYMNANDWGPSENLN" (the letters in the grid sequentially)

# Encryption
def grille_enc (m, k):
    n, holes = k
    assert len(m) == len(holes)
    pass #TODO

# Decryption
def grille_dec (c, k):
    n, holes = k
    assert len(c) == n*n
    pass #TODO

#@# Testing
def grille_test (): #@#
    print("5. Grille cipher") #@#
    k = (5, [(1,0),(4,0),(2,1),(0,2),(4,2),(3,3),(1,4)]) #@#
    m = "MESSAGE" #@#
    #@# Note: we get different ciphers each time.
    test_cipher(grille_enc, grille_dec, k, m) #@#
    test_cipher(grille_enc, grille_dec, k, m) #@#
    test_cipher(grille_enc, grille_dec, k, m) #@#
    k = (3, [(1,0),(0,1),(2,1),(0,2)]) #@#
    test_cipher(grille_enc, grille_dec, k, "TEXT") #@#
    #@# What's wrong with the following grille?
    k = (4, [(0,0),(1,0),(2,0),(3,0),(0,1),(1,1)]) #@#
    test_cipher(grille_enc, grille_dec, k, "SECRET") #@#
    print() #@#
#grille_test()


#@# ---------------------------------------------------------------------------
#@# Exercise 6: Turning grille cipher
#@#
#@# This is a variant of the grille cipher: now the grille is placed not once,
#@# but four times on the grid, each time rotated 90 degrees clockwise.
#@# The grille must be constructed in such a way that each cell will be filled
#@# exactly once during this process, and so no random letters are needed.
#@# This means that the message must be exactly as long as the number of cells.
#@# For example, encrypting "PLAINTEXTMESSAGE"
#@#    # P # L     # # # #     # # T #     S # # #       S P T L
#@#    # # # #  +  N # # T  +  # M # #  +  # # A #  -->  N M A T
#@#    # # A #     # E # #     # # # #     G # # E       G E A E
#@#    # I # #     # # # X     E # S #     # # # #       E I S X
#@# The key and the ciphertext is encoded as in the previous exercise.

#@# a) Write a function that rotates the given grille once
#@# (where k is the grille as (n, [(x1,y1), (x2,y2), ...]), see above).
def rotate_grille (k):
    n, holes = k
    rot_holes = holes #TODO (replace this line with your rotation code)
    # Sort the positions (first by y, then by x)
    rot_holes.sort(key = lambda xy: xy[1]*n + xy[0])
    return (n, rot_holes)
#assert rotate_grille((4, [(1,0),(3,0),(2,2),(1,3)])) == (4, [(0,1),(3,1),(1,2),(3,3)])

#@# b) Implement encryption and decryption with the turning grille.

# Encryption
def tgrille_enc (m, k):
    pass #TODO

# Decryption
def tgrille_dec (c, k):
    pass #TODO

#@# Testing
def tgrille_test (): #@#
    print("6. Turning grille cipher") #@#
    k = (4, [(1,0),(3,0),(2,2),(1,3)]) #@#
    m = "PLAINTEXTMESSAGE" #@#
    c = "SPTLNMATGEAEEISX" #@#
    test_cipher(tgrille_enc, tgrille_dec, k, m, c) #@#
    print() #@#
#tgrille_test()

#@# Bonus exercise:
#@# Construct a valid 6x6 turning grille (manually).