#!/usr/bin/env python2

import argparse

try:
    from Crypto.PublicKey import RSA
    from Crypto.Signature import PKCS1_v1_5
    from Crypto.Hash import SHA256
    from base64 import b64decode, b64encode
except ImportError:
    raise ImportError("Please, install PyCrypto. Possible solution for Ubuntu: `sudo apt-get install python-crypto`")

try:
    from pyasn1.type import univ, namedtype
    from pyasn1.codec.der import encoder
except ImportError:
    raise ImportError("Please, install pyasn1. Possible solution for Ubuntu: `sudo apt-get install python-pyasn1`")


class PaData(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('PaVersion', univ.Integer()),
        namedtype.NamedType('PaFlags', univ.Integer()),
        namedtype.NamedType('PaId', univ.OctetString()),
        namedtype.NamedType('FiveSignature', univ.OctetString()),
    )


class PaSignature(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('PaData', PaData()),
        namedtype.NamedType('Signature', univ.OctetString()),
    )


FLAG_ANDROID_PACKAGE = 1 << 0
FLAG_THIRD_PARTY = 1 << 1


class PaCertificate(object):

    VERSION = 1

    def __init__(self, id_data, five_signature, key, third_party=False, android_package=False):
        self.id_data = id_data
        self.five_signature = five_signature
        self.key = key

        signer = PKCS1_v1_5.new(self.key)

        asn1_pa_data = PaData()
        asn1_pa_data['PaVersion'] = self.VERSION
        asn1_pa_data['PaFlags'] = 0
        asn1_pa_data['PaFlags'] |= FLAG_THIRD_PARTY if third_party else 0
        asn1_pa_data['PaFlags'] |= FLAG_ANDROID_PACKAGE if android_package else 0
        asn1_pa_data['PaId'] = SHA256.new(id_data).digest()
        asn1_pa_data['FiveSignature'] = five_signature

        signature = signer.sign(SHA256.new(encoder.encode(asn1_pa_data)))

        asn1_pa_signature = PaSignature()
        asn1_pa_signature['PaData'] = asn1_pa_data
        asn1_pa_signature['Signature'] = signature

        self._asn1 = asn1_pa_signature

    def to_bytes(self):
        return encoder.encode(self._asn1)


class PaSigner(object):
    def __init__(self, **kwargs):
        self.key = None
        self.key_file = None
        for key, value in kwargs.items():
            setattr(self, key, value)

        if not self.key:
            if not self.key_file:
                raise ValueError("Key not specified")

            with open(self.key_file, 'r') as f:
                self.key = RSA.importKey(f.read())

    def native_sign(self, id_data, base64_five_signature):
        return PaCertificate(id_data, b64decode(base64_five_signature), self.key).to_bytes()


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Sign PA files')
    parser.add_argument('--path', metavar='PATH',  help='path to file on device', required=True)
    parser.add_argument('--five-sig', dest='five_sig', metavar='SIG', help='base64 encoded five signature',
                        required=True)
    parser.add_argument('--key', metavar='FILE', help='path to private key', required=True)
    args = parser.parse_args()

    pa_signer = PaSigner(key_file=args.key)
    pa_certificate = pa_signer.native_sign(args.path, args.five_sig)

    print "%s %s" % (args.path, b64encode(pa_certificate))