from impacket.structure import pack,Structure
from Crypto.Cipher import Blowfish
from Crypto.Hash import MD5

import base64

PKT_NONE          =0 #
PKT_PUBKEY_ENC    =1 # /* public key encrypted packet */
PKT_SIGNATURE     =2 # /* secret key encrypted packet */
PKT_SYMKEY_ENC    =3 # /* session key packet (OpenPGP)*/
PKT_ONEPASS_SIG   =4 # /* one pass sig packet (OpenPGP)*/
PKT_SECRET_KEY    =5 # /* secret key */
PKT_PUBLIC_KEY    =6 # /* public key */
PKT_SECRET_SUBKEY =7 # /* secret subkey (OpenPGP) */
PKT_COMPRESSED    =8 # /* compressed data packet */
PKT_ENCRYPTED     =9 # /* conventional encrypted data */
PKT_MARKER        =10 # /* marker packet (OpenPGP) */
PKT_PLAINTEXT     =11 # /* plaintext data with filename and mode */
PKT_RING_TRUST    =12 # /* keyring trust packet */
PKT_USER_ID       =13 # /* user id packet */
PKT_PUBLIC_SUBKEY =14 # /* public subkey (OpenPGP) */
PKT_OLD_COMMENT   =16 # /* comment packet from an OpenPGP draft */
PKT_ATTRIBUTE     =17 # /* PGP's attribute packet */
PKT_ENCRYPTED_MDC =18 # /* integrity protected encrypted data */
PKT_MDC           =19 # /* manipulation detection code packet */
PKT_COMMENT       =61 # /* new comment packet (private) */
PKT_GPG_CONTROL   =63 # /* internal control packet */

CIPHER_ALGO_NONE	 =0
CIPHER_ALGO_IDEA	 =1
CIPHER_ALGO_3DES	 =2
CIPHER_ALGO_CAST5	 =3
CIPHER_ALGO_BLOWFISH	 =4 # /* blowfish 128 bit key */
# /* 5 & 6 are reserved */
CIPHER_ALGO_AES          =7
CIPHER_ALGO_AES192       =8
CIPHER_ALGO_AES256       =9
CIPHER_ALGO_TWOFISH 	 =10 # /* twofish 256 bit */
CIPHER_ALGO_DUMMY        =110 # /* no encryption at all */

PUBKEY_ALGO_RSA        =1
PUBKEY_ALGO_RSA_E      =2   #  /* RSA encrypt only */
PUBKEY_ALGO_RSA_S      =3   #  /* RSA sign only */
PUBKEY_ALGO_ELGAMAL_E  =16   #  /* encrypt only ElGamal (but not for v3)*/
PUBKEY_ALGO_DSA        =17
PUBKEY_ALGO_ELGAMAL    =20   #  /* sign and encrypt elgamal */

PUBKEY_USAGE_SIG     =1	   # /* key is good for signatures */
PUBKEY_USAGE_ENC     =2	   # /* key is good for encryption */
PUBKEY_USAGE_CERT    =4     # /* key is also good to certify other keys*/
PUBKEY_USAGE_AUTH    =8     # /* key is good for authentication */
PUBKEY_USAGE_UNKNOWN =128   # /* key has an unknown usage bit */

DIGEST_ALGO_MD5       =1
DIGEST_ALGO_SHA1      =2
DIGEST_ALGO_RMD160    =3
# /* 4, 5, 6, and 7 are reserved */
DIGEST_ALGO_SHA256    =8
DIGEST_ALGO_SHA384    =9
DIGEST_ALGO_SHA512    =10
DIGEST_ALGO_SHA224    =11

COMPRESS_ALGO_NONE   =0
COMPRESS_ALGO_ZIP    =1
COMPRESS_ALGO_ZLIB   =2
COMPRESS_ALGO_BZIP2  =3

#define is_RSA(a)     ((a)==PUBKEY_ALGO_RSA || (a)==PUBKEY_ALGO_RSA_E \
#		       || (a)==PUBKEY_ALGO_RSA_S )
#define is_ELGAMAL(a) ((a)==PUBKEY_ALGO_ELGAMAL_E)
#define is_DSA(a)     ((a)==PUBKEY_ALGO_DSA)

def BlowfishPGP_new(key, mode, iv = None):
    return BlowfishPGP(key, mode, iv)

class BlowfishPGP:
    new = BlowfishPGP_new
    def __init__(self, key, mode, iv = None):
        self.mode = mode

        if iv is None:
            iv = '\x00' * Blowfish.block_size

        if mode == Blowfish.MODE_PGP:
            mode = Blowfish.MODE_ECB
            self.unused = 0
            self.iv = iv

        self.cipher = Blowfish.new(key, mode, iv)

    def more_iv(self):
        if not self.unused:
            self.last_iv = self.iv
            self.iv = self.cipher.encrypt(self.iv)
            self.unused = len(self.iv)

    def encrypt(self, data):
        def xor(str1, str2):
            return ''.join(map(lambda x,y:chr(ord(x) ^ ord(y)), str1, str2))

        if not self.mode == Blowfish.MODE_PGP:
            return self.cipher.encrypt(data)

        # implement CFB as implemented by OpenGPG
        def one_block(data):
            # called to encrypt blocksize or less bytes
            self.more_iv()	# unused can't be zero

            nbytes = len(data)
            bytos = min(nbytes, self.unused)

            answer = xor(
                self.iv[-self.unused:len(self.iv)-self.unused+bytos],
                data[:bytos])

            self.iv = '%s%s%s' % (
                self.iv[:-self.unused],
                answer,
                self.iv[8-self.unused+bytos:])

            self.unused -= bytos
            if bytos == nbytes:
                return answer

            # here we are if what was available as unused was not enough for this block

        answer = ''
        for i in range(0,len(data),self.cipher.block_size):
            answer += one_block(data[i:i+self.cipher.block_size])

        return answer

def pack_pkt_header(version, data):
    if not hasattr(data, 'type'):
        return ''
    length = len(data)

    if version or data.type > 15:
        ctb = 0xc0
    else:
        ctb = 0x80

    if not ctb & 0x40:
        # old
        ctb |= (data.type & 0x3f) << 2

        length_len = 0
        if length:
            if length < 256:
                length_len = 0
                l = pack('B', length)
            elif length < 66536:
                length_len = 1
                l = pack('!H', length)
            else:
                length_len = 2
                l = pack('!L', length)

        ctb |= length_len
    else:
        # new_ctb
        ctb |= data.type & 0x3f

        if   length < 192:
            l = pack('B',length)
        elif length < 8384:
            length -= 192
            l = pack('BB', (length/256)+192, length % 256)
        else:
            l = pack('!BL', 0xff, length)

    return '%c%s' % (ctb, l)

class Packet(Structure):
    commonHdr = (
        # ('_hdr_len','_-_hdr','unpack_size()'),
        ('_hdr',':=pack_pkt_header(version, data)'),
        ('version','_'),
        ('data',':'),
    )
    type = 0

    def __init__(self, *args, **kargs):
        Structure.__init__(self, *args, **kargs)

        self['pack_pkt_header'] = pack_pkt_header
        self['version'] = 1

class Plaintext(Structure):
    type = PKT_PLAINTEXT
    structure = (
        ('mode', 'B=0'),
        ('namelen','B-name'),
        ('name',':'),
        ('timestamp','!L=0'),
        ('data',':'))

class SymmetricEncryptionKey(Structure):
    S2K_SIMPLE = 0
    S2K_SALTED = 1
    S2K_ITERATED_SALTED = 3

    type = PKT_SYMKEY_ENC
    structure = (
        ('version',     'B=4'),
        ('cipher_algo', 'B'),
        ('s2kmode',     'B=1'),
        ('hash_algo',   'B'),
        ('salt',        ':'),    # does not exist for S2K_SIMPLE
        ('seskey',      ':'),    # for S2K_INTERATED_SALTED the first byte indicates the number of iterations
    )

def pack_multiple(packets):
    pkt = Packet()
    if type(packets) not in [type(()), type([])]:
        packets = (packets,)

    answer = []
    for each in packets:
        pkt['data'] = each
        answer.append(str(pkt))

    return ''.join(answer)

def pack_pkt_compressed(algo, data):
    if algo == COMPRESS_ALGO_ZLIB:
        import zlib

        return zlib.compress(pack_multiple(data))

class Compressed(Structure):
    type = PKT_COMPRESSED
    structure = (
        ('algorithm', 'B'),
        ('data', '_'),
        ('compressed',':=pack_pkt_compressed(algorithm, data)'),
    )

    def __init__(self, *args, **kargs):
        Structure.__init__(self, *args, **kargs)

        self['pack_pkt_compressed'] = pack_pkt_compressed

def pack_pkt_encrypted(algo, iv, data, mdc):
    global cipher_key
    if algo == CIPHER_ALGO_BLOWFISH:
        cipher = BlowfishPGP(cipher_key, Blowfish.MODE_PGP)

        data = pack_multiple(data)

        return cipher.encrypt('%s%s%s%s' % (iv, iv[-2:], data, mdc))

class Encrypted(Structure):
    type = PKT_ENCRYPTED
    structure = (
        ('algorithm', '_'),	# not packed
        ('data', '_'),	# not packed (only encrypted)
        ('iv','_'),
        ('data', '_'),	# not packed (only encrypted)
        ('encrypted',':=pack_pkt_encrypted(algorithm, data)'),
    )

    def __init__(self, *args, **kargs):
        Structure.__init__(self, *args, **kargs)

        self['pack_pkt_encrypted'] = pack_pkt_encrypted
        self['iv'] = '12345678'
        self['mdc'] = ''

class EncryptedMDC(Encrypted):
    type = PKT_ENCRYPTED_MDC
    structure = (
        ('version','B=1'),
        ('algorithm', '_'),	# not packed
        ('iv','_'),
        ('data', '_'),	# not packed (only encrypted)
        ('mdc','_'),
        ('encrypted',':=pack_pkt_encrypted(algorithm, iv, data, mdc)'),
    )

    def __init__(self, *args, **kargs):
        Encrypted.__init__(self, *args, **kargs)

        self['mdc'] = '12345678901234567890'

class OnePassSignature(Structure):
    type = PKT_ONEPASS_SIG
    structure = (
        ('version','B=3'),
        ('sig_class','B=0'),
        ('digest_algo','B'),
        ('pubkey_algo','B'),
        ('keyid','*!L'),
        ('last','B=1'),
    )

class Raw(Structure):
    structure = (('data',':'),)

class PublicSubkeyComment(Structure):
    type = PKT_PUBLIC_SUBKEY
    structure = (
        ('version','"#'),      # version = '#' is old style comments
        ('data',':'),
    )

class MPINumber(Structure):
    structure = (
        ('_dataLen','_-key','dataLen/8'),
        ('dataLen','!H=len(data)*8'),
        ('data',':'),
    )

class ElGamalPublicKey(Structure):
    structure = (
        ('n1',':'),
        ('n2',':'),
        ('n3',':'),
    )

    def __init__(self, *args, **kargs):
        Structure.__init__(self, *args, **kargs)

        self['n1'] = MPINumber()
        self['n2'] = MPINumber()
        self['n3'] = MPINumber()

class ElGamalSecretKey(Structure):
    structure = (
        ('n1',':'),
    )

    def __init__(self, *args, **kargs):
        Structure.__init__(self, *args, **kargs)

        self['n1'] = MPINumber()

class SecretKeyV4(Structure):
    type = PKT_SECRET_KEY
    structure = (
        ('version','B=4'),       # version in [2,3,4]
        ('timestamp','!L=0'),
        ('algo','B'),            # one of PUBKEY_ALGO_*
        ('publicKey',':'),
        ('protect_algo','B'),
        ('secretKey',':'),
        ('cksum','!H'),
    )


def crc24(octets):
    CRC24_INIT = 0xb704ce
    CRC24_POLY = 0x1864cfb
    crc = CRC24_INIT;

    for s in range(len(octets)):
        crc = crc ^ ( ord(octets[s]) << 16 )
        for i in range(8):
            crc = crc << 1;
            if(crc & 0x1000000):
                crc = crc ^ CRC24_POLY
    crc = crc & 0xffffff
    
    c2 = crc & 0xff
    c1 = (crc >> 8) & 0xff
    c0 = (crc >> 16) & 0xff
    
    return base64.encodestring("%c%c%c" % (c0,c1,c2))


if __name__ == '__main__':
    ff = open('sasa.gpg','wb')

    pkt = Packet()

    pkt['version']  = 0

    """
    pkt['data'] = Plaintext()
    pkt['data']['name'] = 'test.txt'
    pkt['data']['data'] = 'hola loco\n'
    ff.write(str(pkt))
    """

    private_key = 'la clave'
    session_key = '12345678901234567890123'
    salt        = '12345678'

    symkey_enc = SymmetricEncryptionKey()
    symkey_enc['cipher_algo'] = CIPHER_ALGO_BLOWFISH
    symkey_enc['hash_algo']   = DIGEST_ALGO_MD5
    symkey_enc['s2kmode']     = 1
    symkey_enc['salt']        = salt

    userkey_hash = MD5.new()
    userkey_hash.update(salt + private_key)
    userkey_hash = userkey_hash.digest()

    if 1:
        # not working. SymmetrycEncryptionKey can have an encripted session key
        # in this case, the session key is embeded and encripted with user's typed key
        # using cipher_algo
        # the algorithm for data encription is the first byte of the decripted sess key


        print "Session Key's Cipher key: (%d) %r" % (len(userkey_hash), userkey_hash)

        cipher = BlowfishPGP(userkey_hash, Blowfish.MODE_PGP, '\x00'*8)
        symkey_enc['seskey']      = cipher.encrypt('\x04' + session_key)

        print "Encrypted session key: (%d) %r" % (len(symkey_enc['seskey']),symkey_enc['seskey'])

        cipher_key = session_key
    else:
        # SymmetricEncryptionKey can be the algorithm definition for the session key
        # which will be Hash("user typed ley")
        # in this case, cipher_algo is the algorithm for data encription

        symkey_enc['s2kmode']     = 1
        symkey_enc['seskey']      = ''
        cipher_key = userkey_hash

    print "Cipher key: (%d) %r" % (len(cipher_key), cipher_key)

    encrypted = EncryptedMDC()
    encrypted['algorithm'] = CIPHER_ALGO_BLOWFISH
    encrypted['data'] = [Plaintext(), Plaintext()]
    encrypted['data'][0]['name'] = 'crypted.txt'
    encrypted['data'][0]['data'] = 'hola manola\n'
    encrypted['data'][1]['name'] = 'also-crypted.txt'
    encrypted['data'][1]['data'] = 'chau pineda\n'

    compressed = Compressed()
    compressed['algorithm'] = COMPRESS_ALGO_ZLIB
    compressed['data'] = [symkey_enc, encrypted]

    skey = SecretKeyV4()
    skey['algo'] = PUBKEY_ALGO_ELGAMAL
    skey['publicKey'] = ElGamalPublicKey()
    skey['publicKey']['n1']['data'] = '12345678'
    skey['publicKey']['n2']['data'] = 'abcdefgh'
    skey['publicKey']['n3']['data'] = 'terceros'
    skey['protect_algo'] = 0
    skey['secretKey'] = ElGamalSecretKey()
    skey['secretKey']['n1']['data'] = '12345678'
    skey['cksum'] = 0x1234
    pkt['data'] = skey
    ff.write(str(pkt))

    pkt['data'] = compressed
    ff.write(str(pkt))

    """
    pkt['data'] = encrypted
    ff.write(str(pkt))

    pkt['data'] = Plaintext()
    pkt['data']['name'] = 'tost.txt'
    pkt['data']['data'] = 'esto, no lo dije yo, ni en pedo!!! todo lo que sigue, es verso!!!\n'
    ff.write(str(pkt))
    """