From 62be3254033d73b13737d7135f11794ce0cdc103 Mon Sep 17 00:00:00 2001
From: Charles Leifer <coleifer@gmail.com>
Date: Wed, 8 Aug 2018 09:30:25 -0500
Subject: [PATCH] Add option to specify URL prefix for sqlite-web application.

Refs #48
---
 README.md                |  1 +
 sqlite_web/sqlite_web.py | 36 ++++++++++++++++++++++++++++++------
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 87e7f3d3..ca81034b 100644
--- a/README.md
+++ b/README.md
@@ -72,6 +72,7 @@ The following options are available:
 * ``-x``, ``--no-browser``: do not open a web-browser when sqlite-web starts.
 * ``-P``, ``--password``: prompt for password to access sqlite-web.
 * ``-r``, ``--read-only``: open database in read-only mode.
+* ``-u``, ``--url-prefix``: URL prefix for application, e.g. "/sqlite-web".
 
 ### Using docker
 
diff --git a/sqlite_web/sqlite_web.py b/sqlite_web/sqlite_web.py
index 02db0d44..7aba9cd6 100755
--- a/sqlite_web/sqlite_web.py
+++ b/sqlite_web/sqlite_web.py
@@ -193,7 +193,7 @@ def login():
     if request.method == 'POST':
         if request.form.get('password') == app.config['PASSWORD']:
             session['authorized'] = True
-            return redirect(session.get('next_url') or '/')
+            return redirect(session.get('next_url') or url_for('index'))
         flash('The password you entered is incorrect.', 'danger')
     return render_template('login.html')
 
@@ -674,6 +674,22 @@ def _close_db(exc):
     if not dataset._database.is_closed():
         dataset.close()
 
+
+class PrefixMiddleware(object):
+    def __init__(self, app, prefix):
+        self.app = app
+        self.prefix = '/%s' % prefix.strip('/')
+        self.prefix_len = len(self.prefix)
+
+    def __call__(self, environ, start_response):
+        if environ['PATH_INFO'].startswith(self.prefix):
+            environ['PATH_INFO'] = environ['PATH_INFO'][self.prefix_len:]
+            environ['SCRIPT_NAME'] = self.prefix
+            return self.app(environ, start_response)
+        else:
+            start_response('404', [('Content-Type', 'text/plain')])
+            return ['URL does not match application prefix.'.encode()]
+
 #
 # Script options.
 #
@@ -715,6 +731,11 @@ def get_option_parser():
         action='store_true',
         dest='read_only',
         help='Open database in read-only mode.')
+    parser.add_option(
+        '-u',
+        '--url-prefix',
+        dest='url_prefix',
+        help='URL prefix for application.')
     return parser
 
 def die(msg, exit_code=1):
@@ -733,7 +754,7 @@ def _open_tab(url):
     thread.daemon = True
     thread.start()
 
-def install_auth_handler(password):
+def install_auth_handler(password, url_prefix=None):
     app.config['PASSWORD'] = password
 
     @app.before_request
@@ -741,15 +762,15 @@ def check_password():
         if not session.get('authorized') and request.path != '/login/' and \
            not request.path.startswith(('/static/', '/favicon')):
             flash('You must log-in to view the database browser.', 'danger')
-            session['next_url'] = request.path
+            session['next_url'] = request.base_url
             return redirect(url_for('login'))
 
-def initialize_app(filename, read_only=False, password=None):
+def initialize_app(filename, read_only=False, password=None, url_prefix=None):
     global dataset
     global migrator
 
     if password:
-        install_auth_handler(password)
+        install_auth_handler(password, url_prefix)
 
     if read_only:
         if sys.version_info < (3, 4, 0):
@@ -767,6 +788,9 @@ def initialize_app(filename, read_only=False, password=None):
     else:
         dataset = SqliteDataSet('sqlite:///%s' % filename, bare_fields=True)
 
+    if url_prefix:
+        app.wsgi_app = PrefixMiddleware(app.wsgi_app, prefix=url_prefix)
+
     migrator = dataset._migrator
     dataset.close()
 
@@ -788,7 +812,7 @@ def main():
                 break
 
     # Initialize the dataset instance and (optionally) authentication handler.
-    initialize_app(args[0], options.read_only, password)
+    initialize_app(args[0], options.read_only, password, options.url_prefix)
 
     if options.browser:
         open_browser_tab(options.host, options.port)