CTF Multicore bruteforcer for `steghide`

steghide_brute.py

#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
import multiprocessing as mp
import os
import subprocess
import time
import sys
try:
    import click
    from colored import fg, stylize
except ModuleNotFoundError:
    print("pip install colored click")
    sys.exit(-1)


class EndBruteforce(LookupError):
    '''raise this when bruteforce needs to end (== found a passphrase)'''


class Bruteforcer:
    def __init__(self, wordlist, ifile):
        self.__path = wordlist
        self.__ifile = ifile
        init_time = time.time()
        self.__run()
        print("Execution time (s): {}".format(
            stylize(int(time.time() - init_time), fg("blue"))))

    def __chunkify(self, size=1024 * 1024):
        fileEnd = os.path.getsize(self.__path)
        with open(self.__path, "rb") as f:
            chunkEnd = f.tell()
            while True:
                chunkStart = chunkEnd
                f.seek(size, 1)
                f.readline()
                chunkEnd = f.tell()
                yield chunkStart, chunkEnd - chunkStart
                if chunkEnd > fileEnd:
                    break

    def _wrap(self, start, size):
        with open(self.__path, encoding='ISO-8859-1') as f:
            f.seek(start)
            lines = f.read(size).splitlines()
            for line in lines:
                self.__bruteforce(line)

    def __bruteforce(self, word):
        command = "steghide extract -sf {} -p '{}' -v --force".format(
            self.__ifile, word)
        stdoutdata = subprocess.getoutput(command)
        if "writing extracted data to" in stdoutdata:
            print(stylize(stdoutdata, fg("green")))
            print("{}: {}".format(
                stylize("[!!] Passphrase found:", fg("yellow")),
                stylize(word, fg("red"))))
            raise EndBruteforce("Passphrase found!")

    def __run(self):
        with mp.Pool(os.cpu_count()) as p:
            jobs = []
            try:
                for ch_start, ch_size in self.__chunkify():
                    jobs.append(p.apply_async(self._wrap, (ch_start, ch_size)))

                for j in jobs:
                    j.get()
                # If all jobs ends
                print(
                    stylize(
                        "[EE] Could not extract any data with that wordlist",
                        fg("red")))
                sys.exit(-1)
            except EndBruteforce:
                return


@click.command(context_settings={'help_option_names': ['-h', '--help']})
@click.option(
    "--ifile", "-f", nargs=1, help="File to bruteforce", type=click.File("r"))
@click.option(
    "--wordlist", "-w", nargs=1, help="Wordlist to use", type=click.File("r"))
def cmd(ifile, wordlist):
    if ifile and wordlist:
        Bruteforcer(wordlist, ifile)
    else:
        print("steghide_brute.py -h/--help")


if __name__ == "__main__":
    cmd()