#!/usr/bin/env python3
import argparse
from collections import namedtuple

try:
    from Crypto.PublicKey import RSA
    from Crypto.Signature import PKCS1_v1_5
    from Crypto.Hash import SHA256
    from base64 import b64decode
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 CertificateEntry(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('FileName', univ.OctetString()),
        namedtype.NamedType('Certificate', univ.OctetString()),
        namedtype.NamedType('Metadata', univ.OctetString()),
    )

class CertificateSeq(univ.SequenceOf):
    componentType = CertificateEntry()

class CertificateDb(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('CertificateSeq', CertificateSeq()),
        namedtype.NamedType('Version', univ.Integer()),
        namedtype.NamedType('DbMetadata', univ.OctetString()),
    )

class SignedCertificateDb(univ.Sequence):
    componentType = namedtype.NamedTypes(
        namedtype.NamedType('Db', CertificateDb()),
        namedtype.NamedType('Signature', univ.OctetString()),
    )

class Db(object):
    def __init__(self, xattr_file, key_file):
        self.count = 0
        certificates = CertificateSeq()
        XattrRecord = namedtuple('XattrRecord',
                                 ['file_name', 'xattr_name', 'xattr_value'])

        with open(xattr_file, 'rb') as f:
            for line in f.readlines():
                line = line.split()
                try:
                    record = XattrRecord(*line)
                except TypeError:
                    raise ValueError("Invalid format of input file")
                if record.xattr_name != b"user.pa":
                    # Ignore all extended attributes except PROCA
                    continue
                print(record)
                file_name = record.file_name
                certificate = b64decode(record.xattr_value)
                e = CertificateEntry()
                e['FileName'] = file_name
                e['Certificate'] = certificate
                e['Metadata'] = ''
                certificates.append(e)
                self.count += 1

        db = CertificateDb()
        db['CertificateSeq'] = certificates
        db['Version'] = 1
        db['DbMetadata'] = ''

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

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

        signed_db = SignedCertificateDb()
        signed_db['Db'] = db
        signed_db['Signature'] = signature

        self._asn1 = signed_db

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

    def write_to_file(self, file_name):
        with open(file_name, 'wb') as f:
            f.write(self.to_bytes())


def parse_arguments():
    parser = argparse.ArgumentParser(description='Convert PROCA xattrs to db')
    # Format of xattr file: file_name xattr_name base64_encoded_value
    # All xattr_name except user.pa will be ignored
    parser.add_argument('--xattr_file', metavar='XATTR_FILE',
                        help='path to file with xattrs', required=True)
    parser.add_argument('--key', metavar='FILE',
                        help='path to private key', required=True)
    parser.add_argument('-o', '--output',
                        help='path to output file', required=True)
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_arguments()

    proca_db = Db(args.xattr_file, args.key)
    proca_db.write_to_file(args.output)

    print('PROCA: %d entries are stored to DB succesfully' % proca_db.count)