#!/usr/bin/env nix-shell
#!nix-shell -i python -p python3

# NOTE: the colors are used to keep track of which letters have been
#       already substituted, and which are still unknown.

def green(s):
    return f"\033[42m{s}\033[0m"

def red(s):
    return f"\033[41m{s}\033[0m"

def substitute(cipher_text, substitution_map):
    for letter in cipher_text:
        if letter.lower() in substitution_map:
            if letter.isupper():
                x = substitution_map[letter.lower()].upper()
            else:
                x = substitution_map[letter]
            print(green(x), end="")
        else:
            print(red(letter), end="")

with open('message.txt', 'r') as file:
    enc = file.read()

# Here, I will be slowly substituting the letters in the ciphertext
# based on what I believe the words to be. I have commented out the
# invocations of substitute, and commented the newly found letters
# after each invocation.

# Starting off, "wex dqar bg: cbzjZWD{...}" looks like "the flag is: picoCTF{...}"

substitution_map = {
    '0': '0',
    '1': '1',
    '2': '2',
    '3': '3',
    '4': '4',
    '5': '5',
    '6': '6',
    '7': '7',
    '8': '8',
    '9': '9',
    '-': '-',
    '_': '_',
    '\n': '\n',
    ' ': ' ',
    ',': ',',
    '.': '.',
    ':': ':',
    ';': ';',
    '{': '{',
    '}': '}',
    '(': '(',
    ')': ')',
    'w': 't',
    'e': 'h',
    'x': 'e',
    'd': 'f',
    'q': 'l',
    'a': 'a',
    'r': 'g',
    'b': 'i',
    'g': 's',
    'c': 'p',
    'z': 'c',
    'j': 'o',
}

# substitute(enc, substitution_map)

# '(shoft fof captpfe the flag) afe a tspe' -> '(short for capture the flag) are a type'
# 'coipetitiov' -> 'competition'

substitution_map['f'] = 'r'
substitution_map['p'] = 'u'
substitution_map['s'] = 'y'
substitution_map['i'] = 'm'
substitution_map['v'] = 'n'

# substitute(enc, substitution_map)

# 'presentem hith' -> 'presented with'
# 'numter' -> 'number'
# 'security snills' -> 'security skills'

substitution_map['m'] = 'd'
substitution_map['h'] = 'w'
substitution_map['t'] = 'b'
substitution_map['n'] = 'k'

# substitute(enc, substitution_map)

# 'seryice' -> 'service'
# 'FR3LU3NCY' -> 'FR3QU3NCY'

substitution_map['y'] = 'v'
substitution_map['l'] = 'q'

substitute(enc, substitution_map)
