Python Crypto.Cipher.AES 模块，key_size() 实例源码

def test_find_blocksize(self):

    def encryption_oracle(s):
      pt = pkcs7pad(s, AES.block_size)
      if mode == 'ECB':
        return encrypt_ecb(pt, key)
      elif mode == 'CBC':
        return encrypt_cbc(pt, key, iv)

    for mode in ['ECB', 'CBC']:
      for key_size in AES.key_size:
        for _ in xrange(100):
          key = random_bytes(key_size)
          iv = random_bytes(AES.block_size)

          self.assertTrue(AES.block_size == find_blocksize(encryption_oracle))

def test_decrypt_suffix(self):
    '''https://cryptopals.com/sets/2/challenges/14'''

    def encryption_oracle(s):
      data = '%s%s%s' % (pfx, s, sfx)
      return encrypt_ecb(pkcs7pad(data, AES.block_size), key)

    for key_size in AES.key_size:
      for max_size in xrange(0, AES.block_size * 3):
        key = random_bytes(key_size)
        pfx = random_bytes(randint(0, max_size))
        sfx = random_bytes(randint(0, max_size))

        decrypted = decrypt_suffix(encryption_oracle)

        self.assertTrue(decrypted == sfx)

  # }}}

  # CBC {{{
  # TODO add tests with other algorithms (e.g. DES)

def get_password_cry(password, kdf_salt=None, kdf_iters=None):
    """Generate Cry object from password using KDF.

    Optional arguments are not required if you generate your first Cry object.
    If you already used it to encrypt data before, you will have to provide the
    generated arguments (returned by this function) in order to create the
    exact same Cry object.

    Args:
        password (str): User password.
        kdf_salt (:obj:`bytes`, optional): KDF salt. Defaults to None.
        kdf_iters (:obj:`int`, optional): Number of KDF iterations. Defaults to
            None.

    Returns:
        (
            Cry: Cry object ready for encryption use,
            bytes: KDF salt,
            int: KDF iterations
        )

        KDF salt and number of KDF iterations are returned for future reference
        if they were generated and not provided. They will be returned in any
        case.
    """
    key_size = AES.key_size[2] # 256 bit key
    kdf_salt = kdf_salt or os.urandom(config.kdf_salt_size)
    kdf_iters = kdf_iters or randint(*config._kdf_iter_range)
    aes_key = PBKDF2(str(password), kdf_salt, key_size, kdf_iters)
    crypto = Cry(aes_key)
    return crypto, kdf_salt, kdf_iters

def get_rsa_cry(rsa_key, enc_aes=None):
    """Generate Cry object using RSA key.

    Optional arguemnt is not required if you generate your first Cry object. If
    you already used it to encrypt data before, you will have to provide the
    generated argument (returned by this function) in order to create the exact
    same Cry object.

    Args:
        rsa_key (bytes): Public or private RSA key. If public key is used, only
            encryption methods will be available in returned Cry object. The
            RSA key is used to encrypt random generated 256 bit AES key.
        enc_aes (:obj:`bytes`, optional): RSA encrypted AES key. It has to be
            encrypted with the same key provided in rsa_key argument. Defaults
            to None.

    Returns:
        (
            Cry: Cry object ready for encryption use,
            int: RSA key size,
            bytes: encrypted AES key
        )

        RSA key size and encrypted AES key are returned to be stored for future
        reference. If encrypted AES key was not provided, it will be generated,
        but it is returned in any case.
    """
    key_size = AES.key_size[2]
    rsa = RSA.importKey(rsa_key)
    rsa_size = int((rsa.size()+1)/8)
    aes_key = None
    if enc_aes:
        aes_key = rsa.decrypt(enc_aes)
        aes_key = b'\x00' * (AES.block_size - len(aes_key)) + aes_key
    else:
        aes_key = os.urandom(key_size)
        enc_aes = rsa.encrypt(aes_key, 'K')[0]
        enc_aes = b'\x00' * (rsa_size - len(enc_aes)) + enc_aes
    crypto = Cry(aes_key)
    return crypto, rsa_size, enc_aes

def test_cry_password_enc():
    c, kdf_size, kdf_iters = cry.get_password_cry(os.urandom(AES.key_size[2]))
    data = os.urandom(5000)
    assert len(c.enc(data)) % c.vs == 0
    assert data != c.enc(data)[:len(data)]
    assert c.enc(data) != c.enc(data)

def test_cry_password_dec():
    c, kdf_size, kdf_iters = cry.get_password_cry(os.urandom(AES.key_size[2]))
    data = os.urandom(5000)
    dec_data = c.dec(c.enc(data))
    assert data == dec_data[:len(data)]

def test_cry_password_bad_ic():
    c, kdf_size, kdf_iters = cry.get_password_cry(os.urandom(AES.key_size[2]))
    data = os.urandom(5000)
    enc_data = c.enc(data)
    bad_enc_data = enc_data[:1000] + bytes(100) + enc_data[1100:]
    dec_data = c.dec(enc_data)
    error = None
    try:
        bad_dec_data = c.dec(bad_enc_data)
    except Exception as e:
        error = e
    assert type(error) == IntegrityCheckFail

def test_cry_rsa_bad_ic():
    c, ks, ki = cry.get_password_cry(os.urandom(AES.key_size[2]))
    data = os.urandom(5000)
    enc_data = c.enc(data)
    bad_enc_data = enc_data[:1000] + bytes(100) + enc_data[1100:]
    dec_data = c.dec(enc_data)
    error = None
    try:
        bad_dec_data = c.dec(bad_enc_data)
    except Exception as e:
        error = e
    assert type(error) == IntegrityCheckFail


## usecase tests

def test_sizeof_pfxsfx(self):

    def encryption_oracle(s):
      data = '%s%s%s' % (pfx, s, sfx)
      return encrypt_ecb(pkcs7pad(data, AES.block_size), key)

    for key_size in AES.key_size:
      for max_size in xrange(0, AES.block_size * 3):
        key = random_bytes(key_size)
        pfx = random_bytes(randint(0, max_size))
        sfx = random_bytes(randint(0, max_size))

        pfx_size, sfx_size, _ = sizeof_pfxsfx(encryption_oracle, AES.block_size)

        self.assertTrue((pfx_size, sfx_size) == (len(pfx), len(sfx)))

def test_cbc_encrypt_decrypt(self):
    for key_size in AES.key_size:
      for msg_size in xrange(AES.block_size * 3):
        msg = random_bytes(msg_size)
        key = random_bytes(key_size)
        iv = random_bytes(AES.block_size)

        enc = encrypt_cbc(pkcs7pad(msg, AES.block_size), key, iv)
        dec = pkcs7unpad(decrypt_cbc(enc, key))

        self.assertTrue(dec == msg)

def test_ctr_encrypt_decrypt(self):
    '''https://cryptopals.com/sets/3/challenges/18'''

    for key_size in AES.key_size:
      for msg_size in xrange(1, 100):
        key = random_bytes(key_size)
        msg = random_bytes(msg_size)

      self.assertTrue(CTRCipher(key, 0).decrypt(CTRCipher(key, 0).encrypt(msg)) == msg)

def test_rsa_encryptdecrypt(self):
    for key_size in [1024, 2048]: # bigger key sizes take forever
      for msg_size in range(1, 50):
        privkey, pubkey = keygen_rsa(key_size, 0x10001)

        pt = bytes_to_long(random_bytes(msg_size))
        ct = encrypt_rsa(pubkey, pt)

        self.assertTrue(decrypt_rsa(privkey, ct) == pt)

def encrypt(plaintext, key=config.SECRET, key_salt='', no_iv=False):
    """Encrypt shit the right way"""

    # sanitize inputs
    key = SHA256.new(key + key_salt).digest()
    if len(key) not in AES.key_size:
        raise Exception()
    if isinstance(plaintext, unicode):
        plaintext = plaintext.encode('utf-8')

    # pad plaintext using PKCS7 padding scheme
    padlen = AES.block_size - len(plaintext) % AES.block_size
    plaintext += chr(padlen) * padlen

    # generate random initialization vector using CSPRNG
    iv = '\0' * AES.block_size if no_iv else get_random_bytes(AES.block_size)

    # encrypt using AES in CFB mode
    ciphertext = AES.new(key, AES.MODE_CFB, iv).encrypt(plaintext)

    # prepend iv to ciphertext
    if not no_iv:
        ciphertext = iv + ciphertext

    # return ciphertext in hex encoding
    return ciphertext.encode('hex')

def decrypt(ciphertext, key=config.SECRET, key_salt='', no_iv=False):
    """Decrypt shit the right way"""

    # sanitize inputs
    key = SHA256.new(key + key_salt).digest()
    if len(key) not in AES.key_size:
        raise Exception()
    if len(ciphertext) % AES.block_size:
        raise Exception()
    try:
        ciphertext = ciphertext.decode('hex')
    except TypeError:
        log.warning("Ciphertext wasn't given as a hexadecimal string.")

    # split initialization vector and ciphertext
    if no_iv:
        iv = '\0' * AES.block_size
    else:
        iv = ciphertext[:AES.block_size]
        ciphertext = ciphertext[AES.block_size:]

    # decrypt ciphertext using AES in CFB mode
    plaintext = AES.new(key, AES.MODE_CFB, iv).decrypt(ciphertext)

    # validate padding using PKCS7 padding scheme
    padlen = ord(plaintext[-1])
    if padlen < 1 or padlen > AES.block_size:
        raise Exception()
    if plaintext[-padlen:] != chr(padlen) * padlen:
        raise Exception()
    plaintext = plaintext[:-padlen]

    return plaintext

def decrypt(self, encimg, refkey):
        derivative = HKDFv3().deriveSecrets(refkey, binascii.unhexlify(self.cryptKeys), 112)
        parts = ByteUtil.split(derivative, 16, 32)
        iv = parts[0]
        cipherKey = parts[1]
        e_img = encimg[:-10]
        AES.key_size = 128
        cr_obj = AES.new(key=cipherKey, mode=AES.MODE_CBC, IV=iv)
        return cr_obj.decrypt(e_img)