#!/usr/bin/env python3
# LICENSE: BSD-3
# Copyright: Josh Pitts @midnite_runr

import sys
import struct
import shutil
import io
from optparse import OptionParser


def gather_file_info_win(binary):
        """
        Borrowed from BDF...
        I could just skip to certLOC... *shrug*
        """
        flItms = {}
        binary = open(binary, 'rb')
        binary.seek(int('3C', 16))
        flItms['buffer'] = 0
        flItms['JMPtoCodeAddress'] = 0
        flItms['dis_frm_pehdrs_sectble'] = 248
        flItms['pe_header_location'] = struct.unpack('<i', binary.read(4))[0]
        # Start of COFF
        flItms['COFF_Start'] = flItms['pe_header_location'] + 4
        binary.seek(flItms['COFF_Start'])
        flItms['MachineType'] = struct.unpack('<H', binary.read(2))[0]
        binary.seek(flItms['COFF_Start'] + 2, 0)
        flItms['NumberOfSections'] = struct.unpack('<H', binary.read(2))[0]
        flItms['TimeDateStamp'] = struct.unpack('<I', binary.read(4))[0]
        binary.seek(flItms['COFF_Start'] + 16, 0)
        flItms['SizeOfOptionalHeader'] = struct.unpack('<H', binary.read(2))[0]
        flItms['Characteristics'] = struct.unpack('<H', binary.read(2))[0]
        #End of COFF
        flItms['OptionalHeader_start'] = flItms['COFF_Start'] + 20

        #if flItms['SizeOfOptionalHeader']:
            #Begin Standard Fields section of Optional Header
        binary.seek(flItms['OptionalHeader_start'])
        flItms['Magic'] = struct.unpack('<H', binary.read(2))[0]
        flItms['MajorLinkerVersion'] = struct.unpack("!B", binary.read(1))[0]
        flItms['MinorLinkerVersion'] = struct.unpack("!B", binary.read(1))[0]
        flItms['SizeOfCode'] = struct.unpack("<I", binary.read(4))[0]
        flItms['SizeOfInitializedData'] = struct.unpack("<I", binary.read(4))[0]
        flItms['SizeOfUninitializedData'] = struct.unpack("<I",
                                                               binary.read(4))[0]
        flItms['AddressOfEntryPoint'] = struct.unpack('<I', binary.read(4))[0]
        flItms['PatchLocation'] = flItms['AddressOfEntryPoint']
        flItms['BaseOfCode'] = struct.unpack('<I', binary.read(4))[0]
        if flItms['Magic'] != 0x20B:
            flItms['BaseOfData'] = struct.unpack('<I', binary.read(4))[0]
        # End Standard Fields section of Optional Header
        # Begin Windows-Specific Fields of Optional Header
        if flItms['Magic'] == 0x20B:
            flItms['ImageBase'] = struct.unpack('<Q', binary.read(8))[0]
        else:
            flItms['ImageBase'] = struct.unpack('<I', binary.read(4))[0]
        flItms['SectionAlignment'] = struct.unpack('<I', binary.read(4))[0]
        flItms['FileAlignment'] = struct.unpack('<I', binary.read(4))[0]
        flItms['MajorOperatingSystemVersion'] = struct.unpack('<H',
                                                                   binary.read(2))[0]
        flItms['MinorOperatingSystemVersion'] = struct.unpack('<H',
                                                                   binary.read(2))[0]
        flItms['MajorImageVersion'] = struct.unpack('<H', binary.read(2))[0]
        flItms['MinorImageVersion'] = struct.unpack('<H', binary.read(2))[0]
        flItms['MajorSubsystemVersion'] = struct.unpack('<H', binary.read(2))[0]
        flItms['MinorSubsystemVersion'] = struct.unpack('<H', binary.read(2))[0]
        flItms['Win32VersionValue'] = struct.unpack('<I', binary.read(4))[0]
        flItms['SizeOfImageLoc'] = binary.tell()
        flItms['SizeOfImage'] = struct.unpack('<I', binary.read(4))[0]
        flItms['SizeOfHeaders'] = struct.unpack('<I', binary.read(4))[0]
        flItms['CheckSum'] = struct.unpack('<I', binary.read(4))[0]
        flItms['Subsystem'] = struct.unpack('<H', binary.read(2))[0]
        flItms['DllCharacteristics'] = struct.unpack('<H', binary.read(2))[0]
        if flItms['Magic'] == 0x20B:
            flItms['SizeOfStackReserve'] = struct.unpack('<Q', binary.read(8))[0]
            flItms['SizeOfStackCommit'] = struct.unpack('<Q', binary.read(8))[0]
            flItms['SizeOfHeapReserve'] = struct.unpack('<Q', binary.read(8))[0]
            flItms['SizeOfHeapCommit'] = struct.unpack('<Q', binary.read(8))[0]

        else:
            flItms['SizeOfStackReserve'] = struct.unpack('<I', binary.read(4))[0]
            flItms['SizeOfStackCommit'] = struct.unpack('<I', binary.read(4))[0]
            flItms['SizeOfHeapReserve'] = struct.unpack('<I', binary.read(4))[0]
            flItms['SizeOfHeapCommit'] = struct.unpack('<I', binary.read(4))[0]
        flItms['LoaderFlags'] = struct.unpack('<I', binary.read(4))[0]  # zero
        flItms['NumberofRvaAndSizes'] = struct.unpack('<I', binary.read(4))[0]
        # End Windows-Specific Fields of Optional Header
        # Begin Data Directories of Optional Header
        flItms['ExportTableRVA'] = struct.unpack('<I', binary.read(4))[0]
        flItms['ExportTableSize'] = struct.unpack('<I', binary.read(4))[0]
        flItms['ImportTableLOCInPEOptHdrs'] = binary.tell()
        #ImportTable SIZE|LOC
        flItms['ImportTableRVA'] = struct.unpack('<I', binary.read(4))[0]
        flItms['ImportTableSize'] = struct.unpack('<I', binary.read(4))[0]
        flItms['ResourceTable'] = struct.unpack('<Q', binary.read(8))[0]
        flItms['ExceptionTable'] = struct.unpack('<Q', binary.read(8))[0]
        flItms['CertTableLOC'] = binary.tell()
        flItms['CertLOC'] = struct.unpack("<I", binary.read(4))[0]
        flItms['CertSize'] = struct.unpack("<I", binary.read(4))[0]
        binary.close()
        return flItms


def copyCert(exe):
    flItms = gather_file_info_win(exe)

    if flItms['CertLOC'] == 0 or flItms['CertSize'] == 0:
        # not signed
        print("Input file Not signed!")
        sys.exit(-1)

    with open(exe, 'rb') as f:
        f.seek(flItms['CertLOC'], 0)
        cert = f.read(flItms['CertSize'])
    return cert


def writeCert(cert, exe, output):
    flItms = gather_file_info_win(exe)
    
    if not output: 
        output = output = str(exe) + "_signed"

    shutil.copy2(exe, output)
    
    print("Output file: {0}".format(output))

    with open(exe, 'rb') as g:
        with open(output, 'wb') as f:
            f.write(g.read())
            f.seek(0)
            f.seek(flItms['CertTableLOC'], 0)
            f.write(struct.pack("<I", len(open(exe, 'rb').read())))
            f.write(struct.pack("<I", len(cert)))
            f.seek(0, io.SEEK_END)
            f.write(cert)

    print("Signature appended. \nFIN.")


def outputCert(exe, output):
    cert = copyCert(exe)
    if not output:
        output = str(exe) + "_sig"

    print("Output file: {0}".format(output))

    open(output, 'wb').write(cert)

    print("Signature ripped. \nFIN.")


def check_sig(exe):
    flItms = gather_file_info_win(exe)
 
    if flItms['CertLOC'] == 0 or flItms['CertSize'] == 0:
        # not signed
        print("Inputfile Not signed!")
    else:
        print("Inputfile is signed!")


def truncate(exe, output):
    flItms = gather_file_info_win(exe)
 
    if flItms['CertLOC'] == 0 or flItms['CertSize'] == 0:
        # not signed
        print("Inputfile Not signed!")
        sys.exit(-1)
    else:
        print( "Inputfile is signed!")

    if not output:
        output = str(exe) + "_nosig"

    print("Output file: {0}".format(output))

    shutil.copy2(exe, output)

    with open(output, "r+b") as binary:
        print('Overwriting certificate table pointer and truncating binary')
        binary.seek(-flItms['CertSize'], io.SEEK_END)
        binary.truncate()
        binary.seek(flItms['CertTableLOC'], 0)
        binary.write(b"\x00\x00\x00\x00\x00\x00\x00\x00")

    print("Signature removed. \nFIN.")


def signfile(exe, sigfile, output):
    flItms = gather_file_info_win(exe)
    
    cert = open(sigfile, 'rb').read()

    if not output: 
        output = output = str(exe) + "_signed"

    shutil.copy2(exe, output)
    
    print("Output file: {0}".format(output))
    
    with open(exe, 'rb') as g:
        with open(output, 'wb') as f:
            f.write(g.read())
            f.seek(0)
            f.seek(flItms['CertTableLOC'], 0)
            f.write(struct.pack("<I", len(open(exe, 'rb').read())))
            f.write(struct.pack("<I", len(cert)))
            f.seek(0, io.SEEK_END)
            f.write(cert)
    print("Signature appended. \nFIN.")


if __name__ == "__main__":
    usage = 'usage: %prog [options]'
    parser = OptionParser()
    parser.add_option("-i", "--file", dest="inputfile", 
                  help="input file", metavar="FILE")
    parser.add_option('-r', '--rip', dest='ripsig', action='store_true',
                  help='rip signature off inputfile')
    parser.add_option('-a', '--add', dest='addsig', action='store_true',
                  help='add signautre to targetfile')
    parser.add_option('-o', '--output', dest='outputfile',
                  help='output file')
    parser.add_option('-s', '--sig', dest='sigfile',
                  help='binary signature from disk')
    parser.add_option('-t', '--target', dest='targetfile',
                  help='file to append signature to')
    parser.add_option('-c', '--checksig', dest='checksig', action='store_true',
                  help='file to check if signed; does not verify signature')
    parser.add_option('-T', '--truncate', dest="truncate", action='store_true',
                  help='truncate signature (i.e. remove sig)')
    (options, args) = parser.parse_args()
    
    # rip signature
    # inputfile and rip to outputfile
    if options.inputfile and options.ripsig:
        print("Ripping signature to file!")
        outputCert(options.inputfile, options.outputfile)
        sys.exit()    

    # copy from one to another
    # inputfile and rip to targetfile to outputfile    
    if options.inputfile and options.targetfile:
        cert = copyCert(options.inputfile)
        writeCert(cert, options.targetfile, options.outputfile)
        sys.exit()

    # check signature
    # inputfile 
    if options.inputfile and options.checksig:
        check_sig(options.inputfile) 
        sys.exit()

    # add sig to target file
    if options.targetfile and options.sigfile:
        signfile(options.targetfile, options.sigfile, options.outputfile)
        sys.exit()
        
    # truncate
    if options.inputfile and options.truncate:
        truncate(options.inputfile, options.outputfile)
        sys.exit()

    parser.print_help()
    parser.error("You must do something!")