import sys
import os
import pefile
import argparse
import logging

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

class PeModifier:
    def __init__(self, bindata, outputpath):
        self.bindata = bindata
        self.opath = outputpath
        self.pe = pefile.PE(data=self.bindata)
        self.filename = "split.{}.exe"
        self.namecounter = -1

    def genName(self):
        self.namecounter += 1
        return self.filename.format(self.namecounter)

    def writePE(self):
        finalpath = os.path.join(self.opath, self.genName())
        logging.info("File written {}".format(finalpath))
        self.pe.write(finalpath)

    def getSection(self, sectionName):
        for section in self.pe.sections:
            if sectionname in section.Name:
                return section.get_data()
    
    def writeSection(self, sectionname, bindata):
        for section in self.pe.sections:
            if sectionname in section.Name:
                section_file_offset = section.PointerToRawData
                self.pe.set_bytes_at_offset(section_file_offset, bindata)

    def blankSectionData(self):
        for section in self.pe.sections:
            logging.info("blanking section {}".format(section.Name))
            bindata = self.getSection(section.Name)
            bindatalen = len(bindata)
            self.writeSection(section.Name, '\x00'*bindatalen)
            self.writePE()
            self.writeSection(section.Name, bindata)

    def blankSectionDataBytes(self, sectionname, blanklen):
        bindata = self.getSection(sectionname)
        bindatalen = len(bindata)
        logging.info("blanking byte in {}".format(sectionname))
        for i in range(0, bindatalen, blanklen):
            moddata = bindata[0:i]
            moddata += ('\x00' * (bindatalen - i))
            self.writeSection(sectionname, bindata)
            self.writePE()
            self.writeSection(sectionname, bindata)
            moddata = ""
        self.writeSection(sectionname, bindata)

    def enc2bin(filepath):
        memfile = ""
        with open(filepath, "rb") as encfile:
            for line in encfile:
                memfile += line.strip()[::-1]
        return memfile.decode("base64")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="split file in a PE aware manner")
    parser.add_argument("targetfile", type=str, help="target file to split")
    parser.add_argument("outputfolder", type=str, help="output folder for splitted files")
    parser.add_argument("--bytes", type=int, help="amount of bytes to blank")
    parser.add_argument("--section", type=str, help="section to blank bytes")
    myargs = parser.parse_args()

    if myargs.bytes is None:
        myargs.bytes = 1000

    if myargs.section is None:
        myargs.section = ".text"

    bindata = enc2bin(myargs.targetfile)
    logging.info("creating in memory binary data")
    pemod = PeModifier(bindata, myargs.outputfolder)
    logging.info("Blanking sections")
    pemod.blankSectionData()
    logging.info("Blanking bytes in section")
    pemod.blankSectionDataBytes(myargs.section, blanklen=myargs.bytes)

