ChallengeBuild.py

__author__ = '@tomereyz'

import argparse
import os


class ChallengeBuild(object):
    def __init__(self, directory, arch, aslr):
        """
        Scan and build challanges
        :param directory:
        :param is_compile:
        :param is_move_flags:
        """
        if not directory:
            directory = os.getcwd()
        self.directory = directory
        self.sources = []
        self.arch = arch
        self.aslr = True if aslr == 'yes' else False

    def __enter__(self):
        self.generate()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        return

    def generate(self):
        if self.aslr:
            architecture_challenges = "exercises-i386-aslr" if self.arch == 'i386' else 'exercises-arm32v7-aslr'
        else:
            architecture_challenges = "exercises-i386" if self.arch == 'i386' else 'exercises-arm32v7'

        challenges_folders = [os.path.join(self.directory, architecture_challenges, f) for f in
                              os.listdir(os.path.join(self.directory, architecture_challenges)) if
                              os.path.isdir(os.path.join(self.directory, architecture_challenges, f))]

        for challenge in challenges_folders:
            potential_challenge = os.path.join(challenge, os.path.basename(challenge)) + '.c'
            if os.path.isfile(potential_challenge):
                self.sources.append((os.path.basename(potential_challenge)))


def main(arch, aslr, directory=None):
    with ChallengeBuild(directory=directory, arch=arch, aslr=aslr) as handle:
        print handle.sources


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Scan directory and prepares Challanges & Flags')
    parser.add_argument('-d', '--directory', help='main directory', required=False, dest='directory')
    parser.add_argument('-a', '--architecture', help='i386/arm32v7', required=True, dest='arch')

    parser.add_argument('-as', '--aslr', help='yes/no', required=True, dest='aslr')

    main(**vars(parser.parse_args()))