Merge pull request GoogleCloudPlatform#160 from GoogleCloudPlatform/ip

lesv · lesv · commit 10854808a92a · 2016-01-29T15:10:38.000-08:00
Using only the first 2 octets of IP addresses.
diff --git a/managed_vms/cloudsql/main.py b/managed_vms/cloudsql/main.py
@@ -14,6 +14,7 @@
 
 import datetime
 import os
+import socket
 
 from flask import Flask, request
 from flask.ext.sqlalchemy import SQLAlchemy
@@ -22,6 +23,15 @@
 app = Flask(__name__)
 
 
+def is_ipv6(addr):
+    """Checks if a given address is an IPv6 address."""
+    try:
+        socket.inet_pton(socket.AF_INET6, addr)
+        return True
+    except socket.error:
+        return False
+
+
 # [START example]
 # Environment variables are defined in app.yaml.
 app.config['SQLALCHEMY_DATABASE_URI'] = os.environ['SQLALCHEMY_DATABASE_URI']
@@ -41,8 +51,16 @@ def __init__(self, timestamp, user_ip):
 
 @app.route('/')
 def index():
+    user_ip = request.remote_addr
+
+    # Keep only the first two octets of the IP address.
+    if is_ipv6(user_ip):
+        user_ip = ':'.join(user_ip.split(':')[:2])
+    else:
+        user_ip = '.'.join(user_ip.split('.')[:2])
+
     visit = Visit(
-        user_ip=request.remote_addr,
+        user_ip=user_ip,
         timestamp=datetime.datetime.utcnow()
     )
 
diff --git a/managed_vms/datastore/main.py b/managed_vms/datastore/main.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import datetime
+import socket
 
 from flask import Flask, request
 from gcloud import datastore
@@ -21,14 +22,31 @@
 app = Flask(__name__)
 
 
+def is_ipv6(addr):
+    """Checks if a given address is an IPv6 address."""
+    try:
+        socket.inet_pton(socket.AF_INET6, addr)
+        return True
+    except socket.error:
+        return False
+
+
 # [START example]
 @app.route('/')
 def index():
     ds = datastore.Client()
 
+    user_ip = request.remote_addr
+
+    # Keep only the first two octets of the IP address.
+    if is_ipv6(user_ip):
+        user_ip = ':'.join(user_ip.split(':')[:2])
+    else:
+        user_ip = '.'.join(user_ip.split('.')[:2])
+
     entity = datastore.Entity(key=ds.key('visit'))
     entity.update({
-        'user_ip': request.remote_addr,
+        'user_ip': user_ip,
         'timestamp': datetime.datetime.utcnow()
     })
 
