Protect website with a password

pogom/app.py

@@ -45,7 +45,16 @@ def is_authenticated(self):
         else:
             return True
 
+    def is_site_authenticated(self):
+        if self.is_authenticated():
+            return True
+        if config.get('SITE_PASSWORD', None) and not request.cookies.get("siteauth") == config['SITE_AUTH_KEY']:
+            return False
+        return True
+
     def fullmap(self):
+        if not self.is_site_authenticated():
+            return redirect(url_for('login'))
         # if 'search_thread' not in [t.name for t in threading.enumerate()]:
         if (not config.get('GOOGLEMAPS_KEY', None) or
                 not config.get('ACCOUNTS', None)):
@@ -66,8 +75,14 @@ def login(self):
             resp = make_response(redirect(url_for('get_config_site')))
             resp.set_cookie('auth', config['AUTH_KEY'])
             return resp
+        if request.form.get('password', None) == config.get('SITE_PASSWORD', None):
+            resp = make_response(redirect(url_for('fullmap')))
+            resp.set_cookie('siteauth', config['SITE_AUTH_KEY'])
+            return resp;
 
     def heatmap_data(self):
+        if not self.is_site_authenticated():
+            return redirect(url_for('login'))
         return jsonify( Pokemon.get_heat_stats() )
 
     def get_config_site(self):
@@ -78,7 +93,8 @@ def get_config_site(self):
             'config.html',
             gmaps_key=config.get('GOOGLEMAPS_KEY', None),
             accounts=config.get('ACCOUNTS', []),
-            password=config.get('CONFIG_PASSWORD', None))
+            config_password=config.get('CONFIG_PASSWORD', None),
+            site_password=config.get('SITE_PASSWORD', None))
 
     def post_config_site(self):
         if not self.is_authenticated():
@@ -92,6 +108,12 @@ def post_config_site(self):
             config['CONFIG_PASSWORD'] = pw
             config['AUTH_KEY'] = ''.join(random.choice(string.lowercase) for _ in range(32))
 
+        pw = request.form.get('sitePassword', None)
+        pw_changed = (pw != config.get('SITE_PASSWORD', None))
+        if pw_changed:
+            config['SITE_PASSWORD'] = pw
+            config['SITE_AUTH_KEY'] = ''.join(random.choice(string.lowercase) for _ in range(32))
+
         accounts_str = request.form.get('accounts', None)
 
         usernames_before = set([])
@@ -117,7 +139,8 @@ def post_config_site(self):
             'config.html',
             gmaps_key=config.get('GOOGLEMAPS_KEY', None),
             accounts=config.get('ACCOUNTS', []),
-            password=config.get('CONFIG_PASSWORD', None),
+            config_password=config.get('CONFIG_PASSWORD', None),
+            site_password=config.get('SITE_PASSWORD', None),
             alert=True))
         if pw_changed:
             resp.set_cookie('auth', config['AUTH_KEY'])
@@ -137,6 +160,7 @@ def save_config(self):
         with open(config_path, 'w') as f:
             data = {'GOOGLEMAPS_KEY': config['GOOGLEMAPS_KEY'],
                     'CONFIG_PASSWORD': config['CONFIG_PASSWORD'],
+                    'SITE_PASSWORD': config['SITE_PASSWORD'],
                     'SCAN_LOCATIONS': self.scan_config.SCAN_LOCATIONS.values(),
                     'ACCOUNTS': config['ACCOUNTS']}
             f.write(json.dumps(data))
@@ -173,6 +197,8 @@ def map_data(self):
         return jsonify(d)
 
     def cover(self):
+        if not self.is_site_authenticated():
+            return redirect(url_for('login'))
         return jsonify({'cover': self.scan_config.COVER,
                         'scan_locations': self.scan_config.SCAN_LOCATIONS.values()})
 
@@ -208,6 +234,8 @@ def delete_location(self):
         return ('', 204)
 
     def stats(self):
+        if not self.is_site_authenticated():
+            return redirect(url_for('login'))
         stats = Pokemon.get_stats()
         count = sum(p['count'] for p in stats)
         return render_template('stats.html', pokemons=Pokemon.get_stats(), total=count)

runserver.py

@@ -33,12 +33,16 @@ def read_config(scan_config):
 
     config['GOOGLEMAPS_KEY'] = c.get('GOOGLEMAPS_KEY', None)
     config['CONFIG_PASSWORD'] = c.get('CONFIG_PASSWORD', None)
+    config['SITE_PASSWORD'] = c.get('SITE_PASSWORD', None)
     config['ACCOUNTS'] = c.get('ACCOUNTS', [])
     scan_config.update_scan_locations(c.get('SCAN_LOCATIONS', {}))
 
     if config.get('CONFIG_PASSWORD', None):
         config['AUTH_KEY'] = ''.join(random.choice(string.lowercase) for _ in range(32))
 
+    if config.get('SITE_PASSWORD', None):
+        config['SITE_AUTH_KEY'] = ''.join(random.choice(string.lowercase) for _ in range(32))
+
 
 if __name__ == '__main__':
     args = get_args()

templates/config.html

@@ -35,9 +35,14 @@ <h1>PoGoMap Config</h1>
       </div>
       <div class="form-group">
         <label for="configPassword">Config Password</label>
-        <input type="password" class="form-control" id="configPassword" name="configPassword" value="{{ password or ''}}">
+        <input type="password" class="form-control" id="configPassword" name="configPassword" value="{{ config_password or ''}}">
         <span id="configPasswordHelp" class="help-block">Protect this config with a password.</span>
       </div>
+      <div class="form-group">
+        <label for="sitePassword">Site Password</label>
+        <input type="password" class="form-control" id="sitePassword" name="sitePassword" value="{{ site_password or ''}}">
+        <span id="sitePasswordHelp" class="help-block">Protect this website with a password.</span>
+      </div>
       <button type="submit" class="btn btn-primary">Submit</button>
       <a href="../" class="btn btn-default">Back to Map</a>
     </form>