#!/usr/bin/python3
import sys
import shutil
import subprocess
import argparse
import xattr

DIGSIG_ELF_SIG_SIZE = 512

parser = argparse.ArgumentParser()
parser.add_argument('dll', metavar='filename.dll', help='path to dll to sign')
parser.add_argument('-p', '--pgoptions', help='pass options to the privacy guard program')
parser.add_argument('-R', '--replace', help='replace original dll', action='store_true')
args = parser.parse_args()

if not args.dll.endswith(('.dll', '.exe')):
    print("[Error] Must have filename.dll as an argument")
    sys.exit(1)

name = args.dll

try:
    if args.replace:
        new_name = name
    else:
        new_name = name[:-4] + '_signed' + name[-4:]
        shutil.copyfile(name, new_name)

    with open(new_name, mode='ab') as f:
        f.write(b'\x00' * DIGSIG_ELF_SIG_SIZE)

    bsign_args = ['bsign', '--sign', '--xattr', new_name]
    if args.pgoptions is not None:
        if '--batch' in args.pgoptions:
            bsign_args.append('--nopass')
        bsign_args.extend(['-p', args.pgoptions])

    if subprocess.call(bsign_args):
        print("[Error] Calling bsign failure")
        sys.exit(1)

    sig = xattr.getxattr(new_name, 'user.sig')
    xattr.removexattr(new_name, 'user.sig')

    with open(new_name, mode='r+b') as f:
        f.seek(-DIGSIG_ELF_SIG_SIZE, 2)
        f.write(sig)

except Exception as e:
    print(str(e))
    sys.exit(1)