from Cryptodome.Cipher import AES, ChaCha20, Salsa20
from .twofish import Twofish
from Cryptodome.Util import Padding as CryptoPadding
import hashlib
from construct import (
Adapter, BitStruct, BitsSwapped, Container, Flag, Padding, ListContainer, Mapping, GreedyBytes, Int32ul, Switch
)
from lxml import etree
from copy import deepcopy
import base64
import unicodedata
import zlib
import re
import codecs
from io import BytesIO
from collections import OrderedDict
class HeaderChecksumError(Exception):
pass
class CredentialsError(Exception):
pass
class PayloadChecksumError(Exception):
pass
class DynamicDict(Adapter):
"""ListContainer <---> Container
Convenience mapping so we dont have to iterate ListContainer to find
the right item
FIXME: lump kwarg was added to get around the fact that InnerHeader is
not truly a dict. We lump all 'binary' InnerHeaderItems into a single list
"""
def __init__(self, key, subcon, lump=[]):
super(DynamicDict, self).__init__(subcon)
self.key = key
self.lump = lump
# map ListContainer to Container
def _decode(self, obj, context, path):
d = OrderedDict()
for l in self.lump:
d[l] = ListContainer([])
for item in obj:
if item[self.key] in self.lump:
d[item[self.key]].append(item)
else:
d[item[self.key]] = item
return Container(d)
# map Container to ListContainer
def _encode(self, obj, context, path):
l = []
for key in obj:
if key in self.lump:
l += obj[key]
else:
l.append(obj[key])
return ListContainer(l)
def Reparsed(subcon_out):
class Reparsed(Adapter):
"""Bytes <---> Parsed subcon result
Takes in bytes and reparses it with subcon_out"""
def _decode(self, data, con, path):
return subcon_out.parse(data, **con)
def _encode(self, obj, con, path):
return subcon_out.build(obj, **con)
return Reparsed
# is the payload compressed?
CompressionFlags = BitsSwapped(
BitStruct("compression" / Flag, Padding(8 * 4 - 1))
)
# -------------------- Key Computation --------------------
def aes_kdf(key, rounds, key_composite):
"""Set up a context for AES128-ECB encryption to find transformed_key"""
cipher = AES.new(key, AES.MODE_ECB)
# get the number of rounds from the header and transform the key_composite
transformed_key = key_composite
for _ in range(0, rounds):
transformed_key = cipher.encrypt(transformed_key)
return hashlib.sha256(transformed_key).digest()
def compute_key_composite(password=None, keyfile=None):
"""Compute composite key.
Used in header verification and payload decryption."""
# hash the password
if password:
password_composite = hashlib.sha256(password.encode('utf-8')).digest()
else:
password_composite = b''
# hash the keyfile
if keyfile:
# try to read XML keyfile
try:
with open(keyfile, 'r') as f:
tree = etree.parse(f).getroot()
keyfile_composite = base64.b64decode(tree.find('Key/Data').text)
# otherwise, try to read plain keyfile
except (etree.XMLSyntaxError, UnicodeDecodeError):
try:
with open(keyfile, 'rb') as f:
key = f.read()
try:
int(key, 16)
is_hex = True
except ValueError:
is_hex = False
# if the length is 32 bytes we assume it is the key
if len(key) == 32:
keyfile_composite = key
# if the length is 64 bytes we assume the key is hex encoded
elif len(key) == 64 and is_hex:
keyfile_composite = codecs.decode(key, 'hex')
# anything else may be a file to hash for the key
else:
keyfile_composite = hashlib.sha256(key).digest()
except:
raise IOError('Could not read keyfile')
else:
keyfile_composite = b''
# create composite key from password and keyfile composites
return hashlib.sha256(password_composite + keyfile_composite).digest()
def compute_master(context):
"""Computes master key from transformed key and master seed.
Used in payload decryption."""
# combine the transformed key with the header master seed to find the master_key
master_key = hashlib.sha256(
context._.header.value.dynamic_header.master_seed.data +
context.transformed_key).digest()
return master_key
# -------------------- XML Processing --------------------
class XML(Adapter):
"""Bytes <---> lxml etree"""
def _decode(self, data, con, path):
return etree.parse(BytesIO(data))
def _encode(self, tree, con, path):
return etree.tostring(tree)
class UnprotectedStream(Adapter):
"""lxml etree <---> unprotected lxml etree
Iterate etree for Protected elements and decrypt using cipher
provided by get_cipher"""
protected_xpath = '//Value[@Protected=\'True\']'
unprotected_xpath = '//Value[@Protected=\'False\']'
def __init__(self, protected_stream_key, subcon):
super(UnprotectedStream, self).__init__(subcon)
self.protected_stream_key = protected_stream_key
def _decode(self, tree, con, path):
cipher = self.get_cipher(self.protected_stream_key(con))
for elem in tree.xpath(self.protected_xpath):
if elem.text is not None:
result = cipher.decrypt(base64.b64decode(elem.text)).decode('utf-8')
# strip invalid XML characters - https://stackoverflow.com/questions/8733233
result = re.sub(
u'[^\u0020-\uD7FF\u0009\u000A\u000D\uE000-\uFFFD\U00010000-\U0010FFFF]+',
'',
result
)
elem.text = result
elem.attrib['Protected'] = 'False'
return tree
def _encode(self, tree, con, path):
tree_copy = deepcopy(tree)
cipher = self.get_cipher(self.protected_stream_key(con))
for elem in tree_copy.xpath(self.unprotected_xpath):
if elem.text is not None:
elem.text = base64.b64encode(
cipher.encrypt(
elem.text.encode('utf-8')
)
)
elem.attrib['Protected'] = 'True'
return tree
class ARCFourVariantStream(UnprotectedStream):
def get_cipher(self, protected_stream_key):
raise Exception("ARCFourVariant not implemented")
# https://github.com/dlech/KeePass2.x/blob/97141c02733cd3abf8d4dce1187fa7959ded58a8/KeePassLib/Cryptography/CryptoRandomStream.cs#L115-L119
class Salsa20Stream(UnprotectedStream):
def get_cipher(self, protected_stream_key):
key = hashlib.sha256(protected_stream_key).digest()
return Salsa20.new(
key=key,
nonce=b'\xE8\x30\x09\x4B\x97\x20\x5D\x2A'
)
# https://github.com/dlech/KeePass2.x/blob/97141c02733cd3abf8d4dce1187fa7959ded58a8/KeePassLib/Cryptography/CryptoRandomStream.cs#L103-L111
class ChaCha20Stream(UnprotectedStream):
def get_cipher(self, protected_stream_key):
key_hash = hashlib.sha512(protected_stream_key).digest()
key = key_hash[:32]
nonce = key_hash[32:44]
return ChaCha20.new(
key=key,
nonce=nonce
)
def Unprotect(protected_stream_id, protected_stream_key, subcon):
"""Select stream cipher based on protected_stream_id"""
return Switch(
protected_stream_id,
{'arcfourvariant': ARCFourVariantStream(protected_stream_key, subcon),
'salsa20': Salsa20Stream(protected_stream_key, subcon),
'chacha20': ChaCha20Stream(protected_stream_key, subcon),
},
default=subcon
)
# -------------------- Payload Encryption/Decompression --------------------
class Concatenated(Adapter):
"""Data Blocks <---> Bytes"""
def _decode(self, blocks, con, path):
return b''.join([block.block_data for block in blocks])
def _encode(self, payload_data, con, path):
blocks = []
# split payload_data into 1 MB blocks (spec default)
i = 0
while i < len(payload_data):
blocks.append(Container(block_data=payload_data[i:i + 2**20]))
i += 2**20
blocks.append(Container(block_data=b''))
return blocks
class DecryptedPayload(Adapter):
"""Encrypted Bytes <---> Decrypted Bytes"""
def _decode(self, payload_data, con, path):
cipher = self.get_cipher(
con.master_key,
con._.header.value.dynamic_header.encryption_iv.data
)
payload_data = cipher.decrypt(payload_data)
return payload_data
def _encode(self, payload_data, con, path):
payload_data = CryptoPadding.pad(payload_data, 16)
cipher = self.get_cipher(
con.master_key,
con._.header.value.dynamic_header.encryption_iv.data
)
payload_data = cipher.encrypt(payload_data)
return payload_data
class AES256Payload(DecryptedPayload):
def get_cipher(self, master_key, encryption_iv):
return AES.new(master_key, AES.MODE_CBC, encryption_iv)
class ChaCha20Payload(DecryptedPayload):
def get_cipher(self, master_key, encryption_iv):
return ChaCha20.new(key=master_key, nonce=encryption_iv)
class TwoFishPayload(DecryptedPayload):
def get_cipher(self, master_key, encryption_iv):
return Twofish.new(master_key, mode=Twofish.MODE_CBC, IV=encryption_iv)
class Decompressed(Adapter):
"""Compressed Bytes <---> Decompressed Bytes"""
def _decode(self, data, con, path):
return zlib.decompress(data, 16 + 15)
def _encode(self, data, con, path):
compressobj = zlib.compressobj(
6,
zlib.DEFLATED,
16 + 15,
zlib.DEF_MEM_LEVEL,
0
)
data = compressobj.compress(data)
data += compressobj.flush()
return data
# -------------------- Cipher Enums --------------------
# payload encryption method
# https://github.com/keepassxreboot/keepassxc/blob/8324d03f0a015e62b6182843b4478226a5197090/src/format/KeePass2.cpp#L24-L26
CipherId = Mapping(
GreedyBytes,
{'aes256': b'1\xc1\xf2\xe6\xbfqCP\xbeX\x05!j\xfcZ\xff',
'twofish': b'\xadh\xf2\x9fWoK\xb9\xa3j\xd4z\xf9e4l',
'chacha20': b'\xd6\x03\x8a+\x8boL\xb5\xa5$3\x9a1\xdb\xb5\x9a'
}
)
# protected entry encryption method
# https://github.com/dlech/KeePass2.x/blob/149ab342338ffade24b44aaa1fd89f14b64fda09/KeePassLib/Cryptography/CryptoRandomStream.cs#L35
ProtectedStreamId = Mapping(
Int32ul,
{'none': 0,
'arcfourvariant': 1,
'salsa20': 2,
'chacha20': 3,
}
)
