Split file while preserving PE format

pe-aware-split.py

#!/usr/bin/env python
# DiabloHorn https://diablohorn.com
# blank out bytes taking into account the PE file format
# input file: base64 malware.exe | rev > enc.txt
import sys
import os
#pip install pefile
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 secion {}".format(section.Name))
            bindata = self.getSection(section.Name)
            bindatalen = len(bindata)
            self.writeSection(section.Name,'\x00'*bindatalen)
            self.writePE()
            #restores in memory pe file to original state
            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,moddata)
            self.writePE()
            #restores in memory pe file to original state
            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)