#!/usr/local/bin/python

import base64
import cgi, cgitb
import hmac
import math
import md5
import os
import random
import re
import sha
import sys
import time
import urllib

SECRET = "FIXME: put random data here"
DEFAULT_P = 155172898181473697471232257763715539915724801966915404479707795314057629378541917580651227423698188993727816152646631438561595825688188889951272158842675419950341258706556549803580104870537681476726513255747040765857479291291572334510643245094715007229621094194349783925984760375594985848253359305585439638443

def xor(a, b):
    assert len(a) == len(b)
    return reduce(lambda x, y: x+y, [chr(ord(x[0]) ^ ord(x[1])) for x in zip(a, b)])

def btwoc(n):
    assert n > 0
    r = ""
    while n > 0:
        r = chr(n % 0x100) + r
        n /= 0x100
    if ord(r[0]) >= 0x80:
        r = chr(0) + r
    return r

def bigint(s):
    s = base64.decodestring(s)
    f = 1
    r = 0
    for i in range(len(s)-1, -1, -1):
        r += f * ord(s[i])
        f *= 0x100
    return r

def randbytes(bytes):
    try:
        f = file("/dev/urandom")
        buf = f.read(bytes)
        f.close()
    except IOError:
        buf = ''.join([chr(random.randint(0, 255)) for x in range(bytes)])
    return buf

def randint(bits):
    try:
        return random.getrandbits(bits)
    except AttributeError:
        buf = randbytes((bits+7) / 8)
        i = 0
        f = 1
        r = 0
        while bits > 0:
            b = ord(buf[i])
            if bits < 8:
                b &= (1 << bits) - 1
            i += 1
            r += f * b
            f *= 0x100
            bits -= 8
        return r

cgitb.enable()

vars = cgi.FieldStorage()
mode = vars.getfirst("openid.mode")
f = file("/tmp/openid.log", "a")
print >>f, time.time(), os.getenv('REMOTE_ADDR'), mode
print >>f, vars
f.close()
if mode == "associate":
    print "Content-type: text/plain"
    print
    response = {}
    response['assoc_type'] = "HMAC-SHA1"
    response['assoc_handle'] = str(time.time())
    response['expires_in'] = 60
    if vars.has_key("openid.session_type") and vars.getfirst("openid.session_type") == "DH-SHA1":
        response['session_type'] = "DH-SHA1"
        p = DEFAULT_P
        if vars.has_key("openid.dh_modulus"):
            p = bigint(vars.getfirst("openid.dh_modulus"))
        g = 2
        if vars.has_key("openid.dh_gen"):
            g = bigint(vars.getfirst("openid.dh_gen"))
        gx = bigint(vars.getfirst("openid.dh_consumer_public"))
        y = randint(int(math.log(p, 2)-1))
        response['dh_server_public'] = base64.encodestring(btwoc(pow(g, y, p))).replace("\n", "")
        response['enc_mac_key'] = base64.encodestring(xor(sha.new(btwoc(pow(gx, y, p))).digest(), hmac.new(SECRET, response['assoc_handle'], sha).digest())).replace("\n", "")
    else:
        response['session_type'] = ""
        response['mac_key'] = base64.encodestring(hmac.new(SECRET, response['assoc_handle'], sha).digest()).replace("\n", "")
    for x in response.items():
        print "%s:%s" % x
elif mode == "checkid_setup":
    if not 'X_HTTP_AUTHORIZATION' in os.environ or not os.environ['X_HTTP_AUTHORIZATION'].lower().startswith("digest "):
        print "Status: 401 Unauthorized"
        print "WWW-Authenticate: Digest realm=\"hewgill.com/OpenID\", nonce=\"%s\"" % base64.encodestring(randbytes(42)).replace("\n", "")
        print "Content-type: text/html"
        print
        print "<p>Unauthorized</p>"
        sys.exit(0)
    a = os.getenv('X_HTTP_AUTHORIZATION')
    auth = {}
    for m in re.finditer(r"(\w+)=(\"?)(.*?)\2(,|$)", a):
        auth[m.group(1)] = m.group(3)
    A1 = md5.new(":".join((auth['username'], auth['realm'], "FIXME: password"))).hexdigest()
    A2 = md5.new(":".join((os.environ['REQUEST_METHOD'], auth['uri']))).hexdigest()
    if 'qop' in auth:
        response = md5.new(":".join((A1, auth['nonce'], auth['nc'], auth['cnonce'], auth['qop'], A2))).hexdigest()
    else:
        response = md5.new(":".join((A1, auth['nonce'], A2))).hexdigest()
    if auth['username'] != "FIXME: username" or auth['response'] != response:
        print "Status: 401 Unauthorized"
        print "WWW-Authenticate: Digest realm=\"hewgill.com/OpenID\", nonce=\"%s\"" % base64.encodestring(randbytes(42)).replace("\n", "")
        print "Content-type: text/html"
        print
        print "<p>Unauthorized</p>"
        sys.exit(0)
    print "Status: 302 Found"
    print "Location:", vars.getfirst("openid.return_to")+''.join(["&"+x[0]+"="+urllib.quote(x[1]) for x in [
        ("openid.mode", "id_res"),
        ("openid.identity", vars.getfirst("openid.identity")),
        ("openid.assoc_handle", vars.getfirst("openid.assoc_handle")),
        ("openid.return_to", vars.getfirst("openid.return_to")),
        ("openid.signed", "mode,identity,return_to"),
        ("openid.sig", base64.encodestring(hmac.new(hmac.new(SECRET, vars.getfirst("openid.assoc_handle"), sha).digest(), ''.join(["%s:%s\n" % x for x in [("mode","id_res"),("identity",vars.getfirst("openid.identity")),("return_to",vars.getfirst("openid.return_to"))]]), sha).digest()).replace("\n", "")),
    ]])
    print "Content-type: text/html"
    print
    print "<p>Please verify that you are %s:</p>" % vars.getfirst("openid.identity")
    print "<p><form><input name=\"password\" /><input type=\"submit\" /></form></p>"
else:
    print "Content-type: text/html"
    print
    print "<p>This is an openid server.</p>"