#!/usr/bin/env python
# -*- coding: iso-8859-1 -*-

import struct
import math
import hashlib
import hmac  # RFC2104
from Crypto.Cipher import XOR

################ PBKDFv2


class PBKDFv2:

    ################ init
    def __init__(self):

        # length of pseudorandom function: 20 for SHA-1, 16 for MD5
        self.hLen = 20

    ################ makeKey
    def makeKey(self, P, S, c, dkLen, digesttype='sha1'):

        # do some sanity checks
        try:
            str(P)
            str(S)
            int(c)
            float(dkLen)
            int(c)
        except:
            print "P = %s, S = %s, c = %s, dkLen = %s:" % (P, S, c, dkLen)
            raise ValueError("ERROR! Input is not correct!")

        if dkLen > ((2 ^ 32 - 1) * self.hLen):
            maxlength = (2 ^ 32 - 1) * self.hLen
            raise ValueError("ERROR! Key is to large! Maxlength is " . str(maxlength))

        l = math.ceil(dkLen / float(self.hLen))

        T = ""
        for blockindex in range(int(l)):
            T += self.F(P, S, c, blockindex, digesttype)
        DK = T[:dkLen]

        return DK

    ################ F
    def F(self, P, S, c, i, digest):

        # The pseudorandom function, PRF, used is HMAC-SHA1 (rfc2104)
        iteration = 1

        # the first iteration; P is the key, and a concatination of
        # S and blocknumber is the message
        istr = struct.pack(">I", i + 1)
        PRFMaster = hmac.new(P, digestmod=getattr(hashlib, digest))
        PRF = PRFMaster.copy()
        PRF.update(S)
        PRF.update(istr)
        U = PRF.digest()  # the first iteration
        Fbuf = U

        while iteration < c:                  # loop through all iterations
            PRF = PRFMaster.copy()
            PRF.update(U)
            U = PRF.digest()
            Fbuf = self._xor(U, Fbuf)    # XOR this new iteration with the old one
            iteration += 1
        return Fbuf

    ################ xor
    def _xor(self, a, b):
        """Performs XOR on two strings a and b"""

        if len(a) != len(b):
            raise ValueError("ERROR: Strings are of different size! %s %s" % (len(a), len(b)))

        xor = XOR.new(a)
        return xor.encrypt(b)


####################################### MAIN ################################################
pbkdf = PBKDFv2()
# pbkdf.makeKey(password, passwordSalt, passwordIterations, keyBytes, 'sha1')
print "(CPU) PBKDFv2 Python output : "+str(''.join(x.encode('hex') for x in pbkdf.makeKey("password", "salt", 1, 20, 'sha1')))