#!/usr/bin/env python3
# Autodict-QL - Optimal token generation for fuzzing
# Part of AFL++ Project
# Author : Microsvuln - Arash.vre@gmail.com
import string
import os
import binascii
import codecs
import struct
import errno
import argparse
import re
import base64
from binascii import unhexlify


def parse_args():
    parser = argparse.ArgumentParser(
        description=(
            "Helper - Specify input file to analysis and output folder to save corpdirus for constants in the overall project -------  Example usage : python2 thisfile.py outdir o.txt"
        )
    )
    parser.add_argument(
        "corpdir", help="The path to the corpus directory to generate files."
    )
    parser.add_argument(
        "infile",
        help="Specify file output of codeql analysis - ex. ooo-hex.txt, analysis take place on this file, example : python2 thisfile.py outdir out.txt",
    )
    return parser.parse_args()


def ensure_dir(dir):
    try:
        os.makedirs(dir)
    except OSError as e:
        if e.errno == errno.EEXIST:
            # print "[-] Directory exists, specify another directory"
            exit(1)


def do_analysis1(corpdir, infile):
    with open(infile, "rb") as f:
        lines = f.readlines()[1:]
        f.close()
        new_lst = []
        n = 1
        for i, num in enumerate(lines):
            if i != 0:
                new_lst.append(num)
                str1 = str(num)
                print("num is " + str1)
                str1 = str1.rstrip("\n\n")
                # str1 = str1.replace("0x","");
                str1 = str1.replace("|", "")
                str1 = str1.rstrip("\r\n")
                str1 = str1.rstrip("\n")
                str1 = str1.replace(" ", "")
                # str1 = str1.translate(None, string.punctuation)
                translator = str.maketrans("", "", string.punctuation)
                str1 = str1.translate(translator)
                str1 = str1[1:]
                str1 = str1[:-1]
                print("After cleanup : " + str1)
                if (
                    (str1 != "0")
                    and (str1 != "ffffffff")
                    and (str1 != "fffffffe")
                    or (len(str1) == 4)
                    or (len(str1) == 8)
                ):
                    print("first : " + str1)
                    if len(str1) > 8:
                        str1 = str1[:-1]
                    elif len(str1) == 5:
                        str1 = str1 = "0"
                    try:
                        # str1 = str1.decode("hex")
                        with open(corpdir + "/lit-seed{0}".format(n), "w") as file:
                            str1 = str1.replace("0x", "")
                            print(str1)
                            str1 = int(str1, base=16)
                            str1 = str1.to_bytes(4, byteorder="little")
                            file.write(str(str1))
                            file.close()
                            with open(corpdir + "/lit-seed{0}".format(n), "r") as q:
                                a = q.readline()
                                a = a[1:]
                                print(
                                    "AFL++ Autodict-QL by Microsvuln : Writing Token :"
                                    + str(a)
                                )
                                q.close()
                                with open(
                                    corpdir + "/lit-seed{0}".format(n), "w"
                                ) as w1:
                                    w1.write(str(a))
                                    print("Done!")
                                    w1.close()
                    except:
                        print("Error!")
                    n = n + 1


def main():
    args = parse_args()
    ensure_dir(args.corpdir)
    do_analysis1(args.corpdir, args.infile)


if __name__ == "__main__":
    main()