import base64
import dns.resolver
import json
import select
import socket

from werkzeug.exceptions import HTTPException
from werkzeug.routing import Map, Rule
from werkzeug.wrappers import Request, Response

import settings

def wait_on_sockets(socks, timeout):
    rs, _, _ = select.select(socks, [], [], timeout)
    for r in rs:
        data = r.recv(4096)
        if data:
            return data
    return None

# Algorithm borrowed from MIT kerberos code. This probably works or
# something.
def send_request(socks, data):
    delay = 2
    for p in range(3):
        for s in socks:
            # Send the request.
            ret = s.send(data)
            if ret == len(data):
                # Wait for a reply for a second.
                reply = wait_on_sockets(socks, 1)
                if reply is not None:
                    return reply
        # Wait for a reply from anyone.
        reply = wait_on_sockets(socks, delay)
        if reply is not None:
            return reply
        delay *= 2
    return None

class WebKDC(object):

    def __init__(self, realm=settings.REALM):
        self.realm = realm
        self.url_map = Map([
            Rule('/v1/<krb_req_b64>', endpoint='request'),
        ])

    def on_request(self, request, krb_req_b64):
        krb_req = base64.b64decode(krb_req_b64)
        krb_rep = self.send_krb_request(krb_req)

        if krb_rep is None:
            data = { 'status': 'TIMEOUT' }
        else:
            data = {
                'status': 'OK',
                'reply': base64.b64encode(krb_rep)
                }
        return Response(json.dumps(data), mimetype='application/json')

    def send_krb_request(self, krb_req):
        # TODO: Support TCP as well as UDP. I think MIT's KDC only
        # support's UDP though.
        srv_query = '_kerberos._udp.' + self.realm
        srv_records = list(dns.resolver.query(srv_query, 'SRV'))
        srv_records.sort(key = lambda r: r.priority)

        socks = []
        try:
            for r in srv_records:
                host = str(r.target)
                port = int(r.port)
                s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                s.setblocking(0)
                s.connect((host, port))
                socks.append(s)

            return send_request(socks, krb_req)
        finally:
            for s in socks:
                s.close()

    def dispatch_request(self, request):
        adapter = self.url_map.bind_to_environ(request.environ)
        try:
            endpoint, values = adapter.match()
            return getattr(self, 'on_' + endpoint)(request, **values)
        except HTTPException, e:
            return e

    def wsgi_app(self, environ, start_response):
        request = Request(environ)
        response = self.dispatch_request(request)
        return response(environ, start_response)

    def __call__(self, environ, start_response):
        return self.wsgi_app(environ, start_response)


def create_app():
    return WebKDC()


if __name__ == '__main__':
    from werkzeug.serving import run_simple
    app = create_app()
    run_simple('127.0.0.1', 5000, app, use_debugger=True, use_reloader=True)

@@ -68,16 +68,158 @@
 socket%0A%0A
+from pyasn1.codec.der import decoder as der_decoder%0Afrom pyasn1.codec.der import encoder as der_encoder%0Afrom pyasn1.error import PyAsn1Error%0A%0A
 from wer
@@ -344,16 +344,32 @@
 sponse%0A%0A
+import krb_asn1%0A
 import s
@@ -1337,13 +1337,16 @@
 /v1/
-%3Ckrb_
+AS_REQ/%3C
 req_
@@ -1362,23 +1362,22 @@
 dpoint='
-request
+AS_REQ
 '),%0A    
@@ -1396,18 +1396,192 @@
 def 
-on_request
+_error_response(self, e):%0A        data = %7B 'status': 'ERROR',%0A                 'msg': str(e) %7D%0A        return Response(json.dumps(data), mimetype='application/json')%0A%0A    def on_AS_REQ
 (sel
@@ -1592,20 +1592,16 @@
 equest, 
-krb_
 req_b64)
@@ -1614,47 +1614,797 @@
     
-krb_req = base64.b64decode(krb_req_b64)
+try:%0A            req_der = base64.b64decode(req_b64)%0A        except TypeError, e:%0A            return self._error_response(e)%0A%0A        # Make sure we don't send garbage to the KDC. Otherwise it%0A        # doesn't reply and we time out, which is kinda awkward.%0A        try:%0A            as_req, rest = der_decoder.decode(req_der,%0A                                              asn1Spec=krb_asn1.AS_REQ())%0A            if rest:%0A                raise ValueError('Garbage after request')%0A            msg_type = int(as_req.getComponentByName('msg-type'))%0A            if msg_type != krb_asn1.KDC_REQ.msg_type_as:%0A                raise ValueError('Bad msg-type')%0A        except (PyAsn1Error, ValueError), e:%0A            return self._error_response(e)%0A%0A        # Okay, it seems good. Go on and send it.
 %0A   
@@ -2440,24 +2440,43 @@
 request(
-krb
+der_encoder.encode(as
 _req)
+)
 %0A%0A      

Validate the contents of a request before sending it