Commit 21a8be00 authored by JackV's avatar JackV
Browse files

Code refactor, agent and referrer handled in separate functions

parent cbebefe7
......@@ -14,6 +14,11 @@ app.config['REDIS_URL'] = os.environ.get('REDIS_URL', 'redis://127.0.0.1:6379/0'
redis_store = FlaskRedis()
redis_store.init_app(app)
pure_string = re.compile(r'^\w+$')
useragent_match = {'Windows': re.compile('Windows'),
'Mac': re.compile('Machintosh'),
'iOS': re.compile('iPhone'),
'Android': re.compile('Linux; Android'),
'Linux': re.compile('Linux')}
def redis_key_for_user(user: str) -> str:
......@@ -23,39 +28,38 @@ def redis_key_for_user(user: str) -> str:
def redis_key_for_url(url: str) -> str:
return "url.{}".format(url)
def redis_key_for_referrer_count(url: str, referrer: str) -> str:
return "url.{}.referrers.{}.count".format(url, referrer)
def redis_key_for_count(url: str) -> str:
return "count.{}".format(url)
def redis_key_for_useragent(url: str) -> str:
return "useragent.{}"
def redirect_useragent(path: str):
useragent = flask.headers.get('User-Agent')
for agent, match in useragent_match.items():
if match.match(useragent):
redis_store.hincrby(redis_key_for_useragent(path),
agent)
def redirect_referrer(path: str):
pass
def handle_redirect(path):
if pure_string.match(path):
final_url = redis_store.get('url.' + path)
if final_url:
referrer = flask.request.headers.get('referer')
if referrer:
safe_referrer = flask.escape(referrer)
if safe_referrer != referrer:
app.logger.warning('XSS injection attempt from {}'.format(
flask.request.remote_addr))
return flask.abort(400)
if not redis_store.exists(redis_key_for_referrer_count(path,
referrer)):
app.logger.debug('Adding {} as referer for {}'.format(
safe_referrer, path))
redis_store.set(redis_key_for_referrer_count(path,
safe_referrer),
0)
redis_store.incr(redis_key_for_referrer_count(path,
safe_referrer))
app.logger.debug('Increased counter for {}'.format(safe_referrer))
redis_store.incr('count.' + path)
return flask.redirect(final_url)
else:
return flask.abort(404)
final_url = redis_store.get('url.' + path)
if final_url:
redirect_useragent(path)
redirect_referrer(path)
redis_store.incr('count.' + path)
return flask.redirect(final_url)
else:
return flask.abort(404)
......@@ -105,9 +109,6 @@ def del_url():
try:
redis_store.delete(redis_key_for_url(payload['name']))
redis_store.delete(redis_key_for_count(payload['name']))
for key in redis_store.scan_iter('url.{}.*'
.format(payload['name'])):
redis_store.delete(key)
return ''
except KeyError:
return flask.abort(400)
......@@ -118,10 +119,8 @@ def del_url():
@app.route('/api/v1/listurls')
@auth.login_required
def list_urls():
keys = filter(lambda z: not re.search('\\.referrers*', z),
map(lambda x: x.decode(), redis_store.keys('url.*')))
return_dict = {}
for key in keys:
for key in redis_store.keys('url.*'):
app.logger.debug(key)
real_key = key.replace('url.', '', 1)
return_dict[real_key] = {'url': redis_store.get(key).decode(),
......@@ -132,26 +131,7 @@ def list_urls():
@app.route('/api/v1/listreferrers/<shortedurl>')
@auth.login_required
def list_referrers(shortedurl=None):
if shortedurl:
# According to this https://stackoverflow.com/questions/22023538/is-it-possible-to-run-a-string-injection-attack-on-a-redis-query
# string injection shouldn't be possible
keys = map(lambda x: x.decode(),
redis_store.scan_iter('url.{}.referrers.*'
.format(shortedurl)))
return_list = []
for key in keys:
referrer_url = re.sub('url\\.{}\\.referrers\\.|\\.count$'
.format(shortedurl),
'', key)
referrer_url_count = int(redis_store.get(
redis_key_for_referrer_count(shortedurl,
referrer_url)).decode())
app.logger.debug(referrer_url)
return_list.append({'url': referrer_url,
'count': referrer_url_count})
return flask.jsonify(return_list)
else:
return flask.jsonify([])
return ''
@app.route('/<path:path>', strict_slashes=False)
......@@ -161,6 +141,7 @@ def redirect(path):
else:
return flask.abort(404)
from .management_cli import *
if __name__ == "__main__":
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment