#!/usr/bin/env python3

from __future__ import print_function
import sys, hexdump
import math, os
import random, time
from simanneal import Annealer

# requires https://github.com/perrygeo/simanneal

cipher_fname=sys.argv[1]
KEYLEN=5

def xor_strings(s,t):
    # https://en.wikipedia.org/wiki/XOR_cipher#Example_implementation
    """xor two strings together"""
    return b"".join(bytes([a^b]) for a,b in zip(s,t))

def read_file(fname):
    file=open(fname, mode='rb')
    content=file.read()
    file.close()
    return content

def chunks(l, n):
    """divide input buffer by n-len chunks"""
    n = max(1, n)
    return [l[i:i + n] for i in range(0, len(l), n)]

trigrams=[b"the", b"and", b"tha", b"ent", b"ing", b"ion", b"tio", b"for", b"nde", b"has", b"nce", b"edt", b"tis", b"oft", b"sth", b"men"]
digrams=[b"th", b"he", b"in", b"er"]

# additional tuning may be needed:
BONUS_FOR_PRINTABLE=1
BONUS_FOR_LOWERCASE=3
BONUS_FOR_SPACE=8
BONUS_FOR_DIGRAM=8
BONUS_FOR_TRIGRAM=16
PENALTY_FOR_DIGIT=-10

cnt=0
def fitness(state):
    global cnt
    fitness=0
    state_as_string=b"".join(map(lambda x: bytes([x]), state))
    tmp=chunks(cipher_file, KEYLEN)
    plain=b"".join(map(lambda x: xor_strings(x, state_as_string), tmp))
    for p in plain:
        if p in b"qwertyuiopasdfghjklzxcvbnm":
            fitness=fitness+BONUS_FOR_LOWERCASE
        if p in b"0123456789":
            fitness=fitness+PENALTY_FOR_DIGIT
        if (p>=0x20 and p<=0x7E) or p==0xA or p==0xD:
            fitness=fitness+BONUS_FOR_PRINTABLE

    for digram in digrams:
        fitness=fitness + plain.count(digram)*BONUS_FOR_DIGRAM

    for trigram in trigrams:
        fitness=fitness + plain.count(trigram)*BONUS_FOR_TRIGRAM
        
    fitness=fitness + plain.count(b" ")*BONUS_FOR_SPACE

    if (cnt%1000)==0:
        print("\u001B[0;0H", end="")
        possible_key=b"".join(map(lambda x: bytes([x]), state))
        print ("state/key:")
        hexdump.hexdump (possible_key)
        print ("decrypted:")
        tmp=chunks(cipher_file, KEYLEN)
        plain=b"".join(map(lambda x: xor_strings(x, possible_key), tmp))
        hexdump.hexdump (plain)
    cnt+=1

    return fitness

def my_update(step, T, E, acceptance, improvement):
    if acceptance!=None and improvement!=None:
        print('Temperature        Energy    Accept   Improve')
        print("%12.5f %12.2f   %7.2f   %7.2f" % (T, E, acceptance, improvement))

class TestProblem(Annealer):

    def __init__(self, state):
        super(TestProblem, self).__init__(state)  # important!
        self.default_update=my_update

    def move(self):
        # set random byte in state to random byte
        i = random.randint(0, len(self.state) - 1)
        self.state[i]=random.randint(0, 255)

    def energy(self):
        return -fitness(self.state)

cipher_file=read_file(cipher_fname)

print("\u001B[2J", end="") # clear screen

init_state = [0 for i in range(KEYLEN)]

test = TestProblem(init_state)
# increase if you're not satisfied with result:
test.steps = 200000
# since our state is just a list, slice is the fastest way to copy
test.copy_strategy = "slice"
#test.copy_strategy = "deepcopy"
state, e = test.anneal()

possible_key=b"".join(map(lambda x: bytes([x]), state))
print ("state/key:")
hexdump.hexdump (possible_key)

print ("decrypted:")
tmp=chunks(cipher_file, KEYLEN)
plain=b"".join(map(lambda x: xor_strings(x, possible_key), tmp))
hexdump.hexdump (plain)

