aglet.py

import numpy as np
import random
import os
import time

data = [
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0,
    0, 0, 0, 1, 0, 0, 0, 0, 2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 2, 0, 0,
    0, 0, 2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0,
    0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 0,
    3, 0, 0, 0, 0, 4, 0, 0, 0, 0, 4, 0, 0, 0, 0, 4, 0, 0, 0, 0, 4, 0, 0, 0, 0, 4,
    0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0,
    0, 0, 0, 1, 0, 0, 0, 0, 2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 2, 0, 0,
    0, 0, 2, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0,
    0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 0, 4, 0, 0, 0, 0, 4, 0, 0, 0, 0, 4, 0, 0, 0, 0,
    4, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
    0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 2, 0,
    0, 0, 0, 2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 3, 0, 0,
    0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0,
    0, 4, 0, 0, 0, 0, 4, 0, 0, 0, 0, 4, 0, 0, 0, 0, 4]

const_data = [
    13, 14, 0,  2,  1,  2,  0,  2,  4,  7,  17, 18, 12, 13, 0,  1,  1,  2,  2,  3,  4,
    5,  17, 18, 0,  3,  8,  9,  5,  6,  10, 11, 5,  6,  12, 14, 12, 15, 16, 19, 12, 15,
    0,  2,  0,  3,  5,  6,  8,  9,  6,  7,  10, 11, 6,  7,  12, 13, 13, 14, 14, 15, 13,
    14, 17, 18, 13, 14, 0,  1,  1,  2,  2,  3,  1,  2,  4,  5,  1,  2,  6,  7,  6,  7,
    12, 13, 4,  6,  14, 15, 8,  11, 17, 19, 13, 14, 0,  2,  1,  3,  5,  6,  1,  3,  9,
    11, 4,  6,  27, 29, 31, 25, 27, 22, 23, 25, 25, 28, 29, 20, 21, 22, 23, 22, 23, 22,
    23, 24, 16, 18, 24, 25, 21, 23, 29, 30, 21, 23, 32, 34, 29, 30, 20, 21, 16, 18, 22,
    23, 19, 30, 31, 20, 21, 32, 33, 22, 23, 34, 35, 30, 31, 21, 23, 16, 17, 24, 25, 16,
    17, 28, 29, 18, 18, 19, 29, 30, 21, 23, 32, 33, 24, 27, 34, 35, 28, 30, 22, 23, 16,
    17, 25, 27, 44, 46, 40, 42, 44, 46, 32, 34, 44, 46, 36, 37, 46, 37, 39, 32, 33, 37,
    39, 34, 35, 37, 39, 36, 37, 46, 47, 42, 43, 46, 47, 44, 33, 40, 41, 34, 35, 42, 43,
    36, 37, 44, 46, 41, 43, 48, 50, 45, 46, 37, 38, 32, 42, 43, 34, 35, 46, 47, 36, 37,
    48, 49, 40, 41, 50, 51, 45, 47, 36, 38, 32, 33, 40, 41, 32, 33, 41, 42, 36, 39, 44,
    46, 40, 43, 48, 49, 44, 46, 37, 39, 32, 34]

n_maxs = [3, 3, 6, 9, 6, 2, 3, 5, 6, 5, 3, 5, 5, 6, 4]

# Known prefix: BountyCon
prefix = np.array([
    [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,
        0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,
        0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [1, 1, 0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 1, 1, 1,
        0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0],
    [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0,
        0, 0, 1, 1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0],
    [1, 1, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 1, 1,
        0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1, 0, 0, 1, 0, 1, 0],
])


def generate_image(offsets):
    assert len(offsets) == len(n_maxs)

    bars = []

    for i in range(0, len(data), 5):
        bars.append(data[i:i+5])

    idx_init = 0
    for n_max, offset in zip(n_maxs, offsets):
        for n in range(n_max):
            def index(i): return 4 * idx_init + (4 * n + offset + i) % (4 * n_max)
            bars[idx_init + n][0] = const_data[index(0)]
            bars[idx_init + n][2] = const_data[index(1)]
            bars[idx_init + n][1] = const_data[index(2)]
            bars[idx_init + n][3] = const_data[index(3)]
        idx_init += n_max

    bars = np.array(bars)

    x1s = bars[:, [0, 2]]
    x2s = bars[:, [1, 3]]
    ys = bars[:, -1]

    img1 = np.zeros([5, 51])
    img2 = np.zeros([5, 51])

    for i, y in enumerate(ys):
        x1 = x1s[i]
        x2 = x2s[i]
        img1[y, x1[0]:x1[1]] = 1
        img2[y, x2[0]:x2[1]] = 1

    image = np.concatenate([img2, img1], 1)

    score = 0

    # give long bars penalty
    d1 = (x1s[:, 1] - x1s[:, 0])
    d2 = (x2s[:, 1] - x2s[:, 0])
    d = np.concatenate([d1, d2], 0)
    score -= (d * d).sum()

    # force it to match the known char
    score -= ((image[:, :prefix.shape[1]] - prefix) ** 2).sum()

    # force the last char be empty
    score -= image[:, -4:].sum() ** 2

    return image, score


def evolution(n_gen):
    class Individual:
        def __init__(self, chromosome):
            self.chromosome = chromosome
            self.data, self.score = generate_image(chromosome)

        def mutate(self, n=1):
            chromosome = self.chromosome.copy()
            for i in range(n):
                j = random.randrange(len(chromosome))
                chromosome[j] = random.randrange(n_maxs[j])
            return self.__class__(chromosome)

    def compute_mean_score(population):
        return sum(individual.score for individual in population) / len(population)

    chromosome = [0] * 15
    population = [Individual(chromosome.copy()) for _ in range(100)]

    for generation in range(n_gen - 1):
        population.sort(key=lambda individual: individual.score)
        yield population

        n = len(population) // 3
        parents = population[n:]
        population = parents + [random.choice(parents).mutate() for _ in range(n)]

    population.sort(key=lambda individual: individual.score)
    yield population


def main():
    score_history = []

    for population in evolution(100):
        score = sum(individual.score for individual in population) / len(population)
        score_history.append(score)
        print('score:', score)

    import matplotlib.pyplot as plt

    plt.figure()

    plt.subplot(2, 1, 1)
    plt.title('flag')
    flag = np.stack([individual.data for individual in population[-5:]]).mean(0)
    plt.imshow(flag)

    plt.subplot(2, 1, 2)
    plt.title('score')
    plt.plot(range(len(score_history)), score_history)

    plt.show()


main()
# BountyCon{v3rt5_4nd_SDL2}