ha and load balancing support

requirements.txt

@@ -1,2 +1,3 @@
 flask
-pykcs11
+pykcs11
+retrying

setup.py

@@ -9,8 +9,9 @@
 version = '0.0.1'
 
 install_requires = [
-    'flask',
-    'pykcs11'
+      'flask',
+      'pykcs11',
+      'retrying'
 ]
 
 setup(name='pyeleven',

src/pyeleven/__init__.py

@@ -1,48 +1,76 @@
+from base64 import b64decode
 from flask import Flask, request, jsonify
-from .pk11 import pkcs11, intarray2bytes, mechanism, find_key, library
+from .pk11 import pkcs11, load_library, slots_for_label
+from .utils import mechanism, intarray2bytes
 import os
-import sys
 import logging
+from .pool import allocation
+from retrying import retry
+from PyKCS11 import PyKCS11Error
 
 __author__ = 'leifj'
 
 app = Flask(__name__)
 app.debug = True
 app.config.from_pyfile(os.path.join(os.getcwd(), 'config.py'))
-app.secret_key = app.config.get("SECRET_KEY")
-print app.config
+max_retry = app.config.get('MAX_RETRY', 7)
+
+
+def pin():
+    return app.config.get('PKCS11PIN', None)
+
+
+def secret_key():
+    return app.config.get("SECRET_KEY")
+
+
+def library_name():
+    return str(app.config['PKCS11MODULE'])
+
+
+#print app.config
 
 logging.basicConfig(level=logging.DEBUG)
 
 
 @app.route("/info")
 def _info():
-    libn = app.config['PKCS11MODULE']
-    return jsonify(dict(library=libn))
+    return jsonify(dict(library=library_name()))
+
 
+@retry(stop_max_attempt_number=max_retry)
+def _do_sign(label, keyname, mech, data, include_cert=True, require_cert=False):
+    if require_cert:
+        include_cert = True
 
-def _find_slot(label):
-    slots = []
-    lib = library(app.config['PKCS11MODULE'])
-    for slot in lib.getSlotList():
-        slot_info = lib.getSlotInfo(slot)
-        if slot_info.get('label') == label:
-            slots.append(slot)
-    return slots
+    with pkcs11(library_name(), label, pin()) as si:
+        key, cert = si.find_key(keyname, find_cert=include_cert)
+        assert key is not None
+        result = dict(slot=label,signed=intarray2bytes(si.session.sign(key, data, mech)).encode('base64'))
+        if require_cert:
+            assert cert is not None
+        if cert and include_cert:
+            result['cert'] = cert
+        return result
 
 
 @app.route("/<slot_or_label>/<keyname>/sign", methods=['POST'])
 def _sign(slot_or_label, keyname):
 
-    slots = []
-    try:
-        slot = int(slot_or_label)
-        slots = [slot]
-    except ValueError:
-        slots = _find_slot(slot_or_label)
+    msg = request.get_json()
+    if not type(msg) is dict:
+        raise ValueError("request must be a dict")
+
+    msg.setdefault('mech', 'RSAPKCS1')
+    if 'data' not in msg:
+        raise ValueError("missing 'data' in request")
+    data = b64decode(msg['data'])
+    mech = mechanism(msg['mech'])
+    return jsonify(_do_sign(slot_or_label, keyname, mech, data, require_cert=True))
+
 
-    if not slots:
-        raise ValueError("No slot found matching %s" % slot_or_label)
+@app.route("/<slot_or_label>/<keyname>/rawsign", methods=['POST'])
+def _rawsign(slot_or_label, keyname):
 
     msg = request.get_json()
     if not type(msg) is dict:
@@ -51,57 +79,47 @@ def _sign(slot_or_label, keyname):
     msg.setdefault('mech', 'RSAPKCS1')
     if 'data' not in msg:
         raise ValueError("missing 'data' in request")
-    data = msg['data'].decode('base64')
-    libn = app.config['PKCS11MODULE']
+    data = b64decode(msg['data'])
     mech = mechanism(msg['mech'])
-    pin = app.config.get('PKCS11PIN', None)
-    for slot in slots:
-        try:
-            with pkcs11(libn, slot, pin=pin) as si:
-                key, cert = find_key(si, keyname)
-                assert key is not None
-                assert cert is not None
-                return jsonify(dict(slot=slot,
-                                    mech=msg['mech'],
-                                    signed=intarray2bytes(si.session.sign(key, data, mech)).encode('base64')))
-        except Exception, ex:
-            logging.error(ex)
-            with pkcs11(libn, slot, pin=pin) as si:
-                si.exception = ex  # invalidate it
-
-    raise ValueError("Unable to sign using any of the matching slots")
+    return jsonify(_do_sign(slot_or_label, keyname, mech, data, include_cert=False))
 
 
 @app.route("/<slot_or_label>", methods=['GET'])
 def _slot(slot_or_label):
-    slot = -1
-    try:
-        slot = int(slot_or_label)
-    except ValueError:
-        slot = _find_slot(slot_or_label)
+    lib = load_library(library_name())
+    slots = slots_for_label(slot_or_label, lib)
+    result = []
+    for slot in slots:
+        r = dict()
+        try:
+            r['mechanisms'] = lib.getMechanismList(slot)
+        except PyKCS11Error, ex:
+            r['mechanisms'] = {'error': str(ex)}
+        try:
+            r['slot'] = lib.getSlotInfo(slot).to_dict()
+        except PyKCS11Error, ex:
+            r['slot'] = {'error': str(ex)}
+        try:
+            r['token'] = lib.getTokenInfo(slot).to_dict()
+        except PyKCS11Error, ex:
+            r['token'] = {'error': str(ex)}
 
-    lib = library(app.config['PKCS11MODULE'])
-    r = dict()
-    try:
-        r['mechanisms'] = lib.getMechanismList(slot)
-    except:
-        pass
-    try:
-        r['slot'] = lib.getSlotInfo(slot).to_dict()
-    except:
-        pass
-    try:
-        r['token'] = lib.getTokenInfo(slot).to_dict()
-    except:
-        pass
-    return jsonify(r)
+        result.append(r)
+
+    return jsonify(dict(slots=result))
 
 
 @app.route("/", methods=['GET'])
 def _token():
-    lib = library(app.config['PKCS11MODULE'])
+    lib = load_library(library_name())
     r = dict()
+    token_labels = dict()
     r['slots'] = lib.getSlotList()
+    for slot in r['slots']:
+        ti = lib.getTokenInfo(slot)
+        lst = token_labels.setdefault(ti.label.strip(), [])
+        lst.append(slot)
+    r['labels'] = token_labels
     return jsonify(r)