router.py

   1 from flask import Flask, redirect, request, render_template
   2 app = Flask(__name__)
   3 import os
   4 from datetime import datetime
   5 from urllib import urlencode
   6 from urlparse import urlparse, parse_qs
   7
   8 from couchdb.client import Server
   9
  10 google_apps_verify = os.environ.get('GOOGLE_VERIFICATION', None)
  11
  12 dashboards = [
  13               os.environ.get('HEROKU_HOST', 'localhost:5000'),
  14              ]
  15
  16 redirects = {
  17 }
  18
  19 couch_url = os.environ.get('CLOUDANT_URL', 'http://localhost:5984/')
  20
  21 server = Server(couch_url)
  22 db_name = 'router'
  23
  24 def get_db(db_name):
  25     if not db_name in server:
  26         server.create(db_name)
  27
  28     return server[db_name]
  29
  30 db = get_db(db_name)
  31
  32 def add_redirect(source, destination, strip_path=False, strip_query=False,
  33         prefer_destination=False):
  34     if source in db:
  35         doc = db[source]
  36     else:
  37         doc = {}
  38     doc['destination']=destination
  39     doc['strip_path']=strip_path
  40     doc['strip_query']=strip_query
  41     doc['prefer_destination']=prefer_destination
  42
  43     db[source]=doc
  44
  45 def get_redirect(source):
  46     if source in db:
  47         return db[source]['destination']
  48     else:
  49         return
  50
  51 def del_redirect(source):
  52     if source in db:
  53         destination = db[source]['destination']
  54         del(db[source])
  55         return destination
  56
  57 class Router(object):
  58     redirects = {}
  59     failure_template = "snap.html"
  60
  61     def __init__(self, redirects=None):
  62         if redirects:
  63             if not isinstance(redirects, dict):
  64                 raise TypeError('Router expects a dict of redirects')
  65             else:
  66                 self.redirects = redirects
  67
  68         self.db = get_db(db_name)
  69
  70     def create(self, src, dst):
  71         self.db[src] = {'destination': dst, 'created': datetime.now()}
  72
  73     def current_redirects(self):
  74         ret = []
  75
  76         for name in self.db:
  77             destination = self.db[name]['destination']
  78             ret.append({'source': name.decode('idna'), 'destination': destination.decode('idna')})
  79
  80         for name in self.redirects.keys():
  81             ret.append({'source': name.decode('idna'), 'destination': self.redirects['name'].decode('idna')})
  82
  83         return ret
  84
  85     def redirect(self, hostname, uri):
  86         if hostname in self.db and 'destination' in self.db[hostname]:
  87             destination = self.db[hostname]['destination']
  88             in_db = True
  89         elif hostname in self.redirects:
  90             destination = self.redirects[hostname]
  91             in_db = False
  92         else:
  93             return render_template(self.failure_template, hostname=hostname, uri=uri)
  94
  95         urlbits = urlparse(destination)
  96
  97         if not urlbits.netloc:
  98             urlbits = urlparse('http://' + destination)
  99             if not urlbits.netloc:
 100                 return render_template(self.failure_template, hostname=hostname, uri=uri)
 101         scheme = urlbits.scheme or 'http'
 102
 103         destination = scheme + '://' + urlbits.netloc
 104
 105         query = {}
 106
 107         if urlbits.query:
 108             query = parse_qs(urlbits.query)
 109             for key, value in query.items():
 110                 query[key] = value[0]
 111         if request.args and (in_db and not self.db[hostname].get('strip_query')):
 112             prefer_destination = in_db and self.db[hostname].get('prefer_destination')
 113             for key, value in request.args.items():
 114                 if key in query and prefer_destination:
 115                     continue
 116                 query[key] = value
 117         if query:
 118             query_string = '?' + urlencode(query)
 119         else:
 120             query_string = ''
 121
 122         if urlbits.fragment:
 123             fragment = '#' + urlbits.fragment
 124         else:
 125             fragment = ''
 126
 127         if in_db and self.db[hostname].get('strip_path'):
 128             path = urlbits.path
 129         else:
 130             path = '/' + '/'.join([
 131                                    urlbits.path.strip('/'),
 132                                    uri.lstrip('/')
 133                                   ])
 134
 135         destination += path + query_string + fragment
 136
 137         if os.environ.get('DEBUG_REDIRECT') and True or False:
 138             return destination
 139
 140         return redirect(destination)
 141
 142 rtr = Router(redirects)
 143
 144 class Dash(object):
 145     urls = []
 146     template = 'dash.html'
 147
 148     def __init__(self, urls=None):
 149         if urls:
 150             if not isinstance(urls, list):
 151                 raise TypeError('Dash expects a list of urls')
 152             else:
 153                 self.urls = urls
 154
 155     def is_dash(self, url):
 156         return url in self.urls
 157
 158     def render(self, uri):
 159         return render_template(self.template, redirects=rtr.current_redirects())
 160
 161     def __call__(self, uri):
 162         return self.render(uri)
 163
 164 dash = Dash(dashboards)
 165
 166 @app.route('/_dash/')
 167 @app.route('/_dash/<path:uri>')
 168 def show_dash(uri=None):
 169     return dash(uri)
 170
 171 @app.route('/')
 172 @app.route('/<path:uri>')
 173 def uri_router(uri=''):
 174     if google_apps_verify and uri == 'google'+google_apps_verify+'.html':
 175         return 'google-site-verification: google'+google_apps_verify+'.html'
 176
 177     hostname = request.host
 178
 179     if dash.is_dash(hostname):
 180         return dash(uri)
 181
 182     return rtr.redirect(hostname, uri)
 183
 184 if __name__ == '__main__':
 185     port = int(os.environ.get('PORT', 5000))
 186     app.debug = os.environ.get('DEBUG') and True or False
 187     app.run('0.0.0.0', port=port)