From 663eec021b41482ea6171582fa35951b92cfdeb6 Mon Sep 17 00:00:00 2001
From: jrconlin <jconlin+git@mozilla.com>
Date: Tue, 28 Jun 2016 14:17:08 -0700
Subject: [PATCH] feat: add AMI instance ID to the logged information

Introduces new shared option "--no_aws" which will prevent calling
AWS meta-information. This option defaults to "False" since autopush
generally always runs on AWS.

No breaking changes.

Closes #483
---
 autopush/endpoint.py        |  1 +
 autopush/main.py            |  9 +++++++++
 autopush/settings.py        |  3 +++
 autopush/tests/test_main.py | 11 +++++++++++
 autopush/utils.py           | 14 ++++++++++++++
 configs/autopush_shared.ini |  6 ++++++
 6 files changed, 44 insertions(+)

diff --git a/autopush/endpoint.py b/autopush/endpoint.py
index 144b398a..36901e38 100644
--- a/autopush/endpoint.py
+++ b/autopush/endpoint.py
@@ -180,6 +180,7 @@ def chid(self, value):
     def _init_info(self):
         """Returns a dict of additional client data"""
         return {
+            "ami_id": self.ap_settings.ami_id,
             "request_id": self.request_id,
             "user_agent": self.request.headers.get("user-agent", ""),
             "remote_ip": self.request.headers.get("x-forwarded-for",
diff --git a/autopush/main.py b/autopush/main.py
index 438d58fd..4f912e48 100644
--- a/autopush/main.py
+++ b/autopush/main.py
@@ -11,6 +11,7 @@
 from twisted.web.server import Site
 
 import autopush.db as db
+import autopush.utils as utils
 from autopush.endpoint import (
     EndpointHandler,
     MessageHandler,
@@ -125,6 +126,8 @@ def add_shared_args(parser):
                         type=int, default=8082, env_var="ENDPOINT_PORT")
     parser.add_argument('--human_logs', help="Enable human readable logs",
                         action="store_true", default=False)
+    parser.add_argument('--no_aws', help="Skip AWS meta information checks",
+                        action="store_true", default=False)
     # No ENV because this is for humans
     add_external_router_args(parser)
     obsolete_args(parser)
@@ -317,6 +320,11 @@ def make_settings(args, **kwargs):
                               "senderIDs": senderIDs,
                               "senderid_list": list}
 
+    ami_id = None
+    # Not a fan of double negatives, but this makes more understandable args
+    if not args.no_aws:
+        ami_id = utils.get_amid()
+
     return AutopushSettings(
         crypto_key=args.crypto_key,
         datadog_api_key=args.datadog_api_key,
@@ -337,6 +345,7 @@ def make_settings(args, **kwargs):
         router_write_throughput=args.router_write_throughput,
         resolve_hostname=args.resolve_hostname,
         wake_timeout=args.wake_timeout,
+        ami_id=ami_id,
         **kwargs
     )
 
diff --git a/autopush/settings.py b/autopush/settings.py
index e5876fe6..06528455 100644
--- a/autopush/settings.py
+++ b/autopush/settings.py
@@ -91,6 +91,7 @@ def __init__(self,
                  hello_timeout=0,
                  bear_hash_key=None,
                  preflight_uaid="deadbeef00000000deadbeef000000000",
+                 ami_id=None,
                  ):
         """Initialize the Settings object
 
@@ -203,6 +204,8 @@ def __init__(self,
 
         self.hello_timeout = hello_timeout
 
+        self.ami_id = ami_id
+
     @property
     def message(self):
         """Property that access the current message table"""
diff --git a/autopush/tests/test_main.py b/autopush/tests/test_main.py
index 77dea8ac..46383332 100644
--- a/autopush/tests/test_main.py
+++ b/autopush/tests/test_main.py
@@ -208,6 +208,7 @@ class test_arg:
         s3_bucket = "none"
         key_hash = "supersikkret"
         senderid_expry = 0
+        no_aws = True
 
     def setUp(self):
         mock_s3().start()
@@ -270,3 +271,13 @@ def test_gcm_start(self, fsi):
             """--senderid_list={"123":{"auth":"abcd"}}""",
             "--s3_bucket=none",
         ])
+
+    @patch("requests.get")
+    def test_aws_ami_id(self, request_mock):
+        class m_reply:
+            content = "ami_123"
+
+        request_mock.return_value = m_reply
+        self.test_arg.no_aws = False
+        ap = make_settings(self.test_arg)
+        eq_(ap.ami_id, "ami_123")
diff --git a/autopush/utils.py b/autopush/utils.py
index aaf0979a..6254b78f 100644
--- a/autopush/utils.py
+++ b/autopush/utils.py
@@ -7,6 +7,7 @@
 import uuid
 
 import ecdsa
+import requests
 from jose import jws
 
 from twisted.logger import Logger
@@ -83,6 +84,19 @@ def base64url_decode(string):
     return base64.urlsafe_b64decode(repad(string))
 
 
+def get_amid():
+    """Fetch the AMI instance ID
+
+    """
+    try:
+        resp = requests.get(
+            "http://169.254.169.254/latest/meta-data/ami-id",
+            timeout=1)
+        return resp.content
+    except:
+        return "Unknown"
+
+
 def decipher_public_key(key_data):
     """A public key may come in several flavors. Attempt to extract the
     valid key bits from keys doing minimal validation checks.
diff --git a/configs/autopush_shared.ini b/configs/autopush_shared.ini
index 36f0d138..696d4383 100644
--- a/configs/autopush_shared.ini
+++ b/configs/autopush_shared.ini
@@ -108,3 +108,9 @@ endpoint_port = 8082
 ; {"12345": {"auth": "abcd_efg"}, "01357": {"auth": "ZYX=abc"}}
 #senderid_list =
 
+; Perform AWS specific information (like fetch the AMI ID from the meta-data
+; server
+;
+; Uncomment to disable AWS meta checks.
+#no_aws
+