#!/usr/bin/python

import os
import sys
import hmac
import cPickle
import sha
import time
import base64
from invirt.vnc import getTokenKey
from invirt.config import structs as config

def getAuthToken(username, machine, lifetime=5*60):
    data = {}
    data['user'] = username
    data['machine'] = machine
    data['expires'] = time.time() + lifetime
    data['connect_host'] = config.web.hostname
    try:
        data['connect_port'] = 10003 + [h.hostname for h in
                                        config.hosts].index(os.uname()[1])
    except:
        data['connect_port'] = 5900
    pickled_data = cPickle.dumps(data)
    m = hmac.new(getTokenKey(), digestmod=sha)
    m.update(pickled_data)
    token = {'data': pickled_data, 'digest': m.digest()}
    token = cPickle.dumps(token)
    token = base64.urlsafe_b64encode(token)
    return token

def main():
    try:
        username = os.environ['REMOTE_USER']
    except KeyError:
        username = None
    machine = sys.argv[1]
    print getAuthToken(username, machine)

if __name__ == '__main__':
    main()
