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

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 file', action='store_true')
parser.add_argument('-w', '--show-info', help='describe signature information', action='store_true')
args = parser.parse_args()

name = args.dll

def show_info():
    f = tempfile.NamedTemporaryFile()
    shutil.copyfile(args.dll, f.name)
    f.seek(-DIGSIG_ELF_SIG_SIZE, 2)
    sig = f.read(DIGSIG_ELF_SIG_SIZE)
    f.seek(-DIGSIG_ELF_SIG_SIZE, 2)
    f.write(b'\x00' * DIGSIG_ELF_SIG_SIZE)
    f.flush()
    xattr.setxattr(f.name, 'user.sig', sig)
    bsign_args = ['bsign', '-w', f.name]
    bsign = subprocess.run(bsign_args, stdout=subprocess.PIPE)
    print(bsign.stdout.decode()
        .replace("xattr ", "dll ")
        .replace(f.name, args.dll))
    return bsign.returncode

try:
    if args.show_info:
        sys.exit(show_info())

    if args.replace:
        new_name = name
    else:
        if not name.endswith(('.dll', '.exe')):
            print("[Error] Must have filename.dll as an argument")
            sys.exit(1)
        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)