Return forwarded IP on prod

1 files changed, 8 insertions, 1 deletions

diff --git a/app.py b/app.py
index 3f67127..805ca8a 100644
--- a/app.py
+++ b/app.py
@@ -1,7 +1,14 @@
 from flask import Flask, request
+from os import environ
 
 app = Flask(__name__)
 
+def get_ip(request):
+    if 'PROD' in environ:
+        return request.headers.get('X-Forwarded-For')
+    else:
+        return request.remote_addr
+
 @app.route('/')
 def index():
-    return request.remote_addr
+    return get_ip(request)