Keepass Protocol.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import base64\n",
    "import itertools\n",
    "import json\n",
    "import os\n",
    "import yaml\n",
    "\n",
    "import requests\n",
    "\n",
    "from cryptography.hazmat.primitives import padding\n",
    "from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes\n",
    "from cryptography.hazmat.backends import default_backend\n",
    "import xdg.BaseDirectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def pad(data, size=128): # size is in bits; 16 bytes = 128 bits\n",
    "    padder = padding.PKCS7(size).padder()\n",
    "    padded_data = padder.update(data)\n",
    "    return padded_data + padder.finalize()\n",
    "\n",
    "def unpad(padded_data, size=128):\n",
    "    unpadder = padding.PKCS7(size).unpadder()\n",
    "    data = unpadder.update(padded_data)\n",
    "    return data + unpadder.finalize()\n",
    "\n",
    "def getCipher(key, iv):\n",
    "    backend = default_backend()\n",
    "    return Cipher(\n",
    "        algorithms.AES(base64.b64decode(key)),\n",
    "        modes.CBC(base64.b64decode(iv)),\n",
    "        backend\n",
    "    )\n",
    "\n",
    "def encrypt(data, key, iv):\n",
    "    cipher = getCipher(key, iv)\n",
    "    encryptor = cipher.encryptor()\n",
    "    p = pad(data)\n",
    "    res = encryptor.update(p) + encryptor.finalize()\n",
    "    return base64.b64encode(res)\n",
    "    \n",
    "def decrypt(data, key, iv):\n",
    "    cipher = getCipher(key, iv)\n",
    "    decryptor = cipher.decryptor()\n",
    "    padded_data = decryptor.update(base64.b64decode(data)) + decryptor.finalize()\n",
    "    return unpad(padded_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "key_size = 32 # this is in bytes\n",
    "iv_size = 16\n",
    "\n",
    "def getRandomBytes(size):\n",
    "    return base64.b64encode(os.urandom(size))\n",
    "\n",
    "def getRandomKey():\n",
    "    return getRandomBytes(key_size)\n",
    "\n",
    "def getRandomIV():\n",
    "    return getRandomBytes(iv_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def getVerifier(iv, key):\n",
    "    return encrypt(iv, key, iv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def merge(d1, d2):\n",
    "    return dict(itertools.chain(d1.iteritems(), d2.iteritems()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def jsonMap(fn, json_obj):\n",
    "    def _fn(v):\n",
    "        if v is None:\n",
    "            return None\n",
    "        elif isinstance(v, dict):\n",
    "            return jsonMap(fn, v)\n",
    "        elif isinstance(v, list):\n",
    "            return map(_fn, v)\n",
    "        else:\n",
    "            return fn(v)\n",
    "    if isinstance(json_obj, dict):\n",
    "        return {k: _fn(v) for k, v in json_obj.iteritems()}\n",
    "    else:\n",
    "        return _fn(json_obj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def encryptDict(dct, key, iv):\n",
    "    def _encrypt(v):\n",
    "        return encrypt(v, key, iv)\n",
    "    return jsonMap(_encrypt, dct)\n",
    "\n",
    "def decryptDict(encrypted_dict, key, iv):\n",
    "    def _decrypt(v):\n",
    "        return decrypt(v, key, iv)\n",
    "    return jsonMap(_decrypt, encrypted_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def checkVerifier(key, iv, verifier):\n",
    "    return verifier == encrypt(iv, key, iv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def convertToStr(input_dict):\n",
    "    return jsonMap(str, input_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "KEEPASS_URL = 'http://localhost:19455/'\n",
    "_output = None\n",
    "\n",
    "def makeRequest(key, input_data, id_, standard_data=None, iv=None):\n",
    "    global _output\n",
    "    _output = None\n",
    "    # standard_data can be set to {} so need to explicitly check that it is equal to None\n",
    "    if standard_data is None:\n",
    "        iv = iv or getRandomIV()\n",
    "        standard_data = {\n",
    "            'Id': id_, \n",
    "            'Nonce': iv, \n",
    "            'Verifier': getVerifier(iv, key)\n",
    "        }\n",
    "    data = merge(standard_data, input_data)\n",
    "    response = requests.post(KEEPASS_URL, json=data)\n",
    "    if response.status_code != 200:\n",
    "        raise Exception('Damn')\n",
    "    output = convertToStr(response.json())\n",
    "    if output['Success']:\n",
    "        if checkVerifier(key, output['Nonce'], output['Verifier']):\n",
    "            return output\n",
    "        else:\n",
    "            print 'Failed to verify'\n",
    "    else:\n",
    "        print 'keepass response was not successful'\n",
    "    # For debugging\n",
    "    _output = output\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def associate():\n",
    "    key = getRandomKey()\n",
    "    input_data = {\n",
    "        'RequestType': 'associate',\n",
    "        'Key': key\n",
    "    }\n",
    "    output = makeRequest(key, input_data, None, {})\n",
    "    if output:\n",
    "        return key, output['Id']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def testAssociate(id_, key):\n",
    "    input_data = {\n",
    "        'RequestType': 'test-associate',\n",
    "        'TriggerUnlock': True,\n",
    "    }\n",
    "    return makeRequest(key, input_data, id_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def getLogins(url, id_, key):\n",
    "    iv = getRandomIV()\n",
    "    input_data = {\n",
    "        'RequestType': 'get-logins',\n",
    "        'Url': encrypt(url, key, iv)\n",
    "    }\n",
    "    output = makeRequest(key, input_data, id_, iv=iv)\n",
    "    if output:\n",
    "        return [decryptDict(entry, key, output['Nonce']) for entry in output['Entries']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def getAndSaveNewAssociation(config_path):\n",
    "    key, id_ = associate()\n",
    "    with open(config_path, 'w') as fout:\n",
    "        fout.write(yaml.safe_dump({'key': key, 'id': id_}, default_flow_style=False))\n",
    "    return key, id_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "config_dir = xdg.BaseDirectory.save_config_path('keepass_test')\n",
    "config_path = os.path.join(config_dir, 'config.yml')\n",
    "if os.path.exists(config_path):\n",
    "    with open(config_path) as fin:\n",
    "        config = yaml.safe_load(fin)\n",
    "    id_ = config['id']\n",
    "    key = config['key']\n",
    "    if not testAssociate(id_, key):\n",
    "        print \"Previous association failed. Loading new association\"\n",
    "        key, id_ = getAndSaveNewAssociation(config_path)\n",
    "else:\n",
    "    print \"No previous association. Loading new association\"\n",
    "    key, id_ = getAndSaveNewAssociation(config_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "logins = getLogins('https://amazon.com', id_, key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'b4dcc4713e72de19750e7fb3c0a10363'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logins[0]['Uuid']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}