Merge pull request RIPE-NCC#127 from danielquinn/issue/122

Issue/122

ripe/atlas/tools/commands/base.py

@@ -1,11 +1,13 @@
 import argparse
+import os
 import re
 import six
 import sys
 
 from datetime import datetime
 
 from ..helpers.colours import colourise
+from ..version import __version__
 
 
 class RipeHelpFormatter(argparse.RawTextHelpFormatter):
@@ -29,6 +31,7 @@ def __init__(self, *args, **kwargs):
             description=self.DESCRIPTION,
             prog="ripe-atlas {}".format(self.NAME)
         )
+        self.user_agent = self._get_user_agent()
 
     def init_args(self, args=None):
         """
@@ -74,6 +77,23 @@ class to manipulate the arguments before being parsed.  The common
     def ok(self, message):
         sys.stdout.write("\n{}\n\n".format(colourise(message, "green")))
 
+    @staticmethod
+    def _get_user_agent():
+        """
+        Allow packagers to change the user-agent to whatever they like by
+        placing a file called `user-agent` into the `tools` directory.  If no
+        file is found, we go with a sensible default + the version.
+        """
+
+        try:
+            custom = os.path.join(os.path.dirname(__file__), "..", "user-agent")
+            with open(custom) as f:
+                return f.readline().strip()[:128]
+        except IOError:
+            pass  # We go with the default
+
+        return "RIPE Atlas Tools (Magellan) {}".format(__version__)
+
 
 class TabularFieldsMixin(object):
     """

ripe/atlas/tools/commands/measure/base.py

@@ -238,6 +238,7 @@ def create(self):
         return AtlasCreateRequest(
             server=conf["ripe-ncc"]["endpoint"].replace("https://", ""),
             key=self.arguments.auth,
+            user_agent=self.user_agent,
             measurements=[creation_class(**self._get_measurement_kwargs())],
             sources=[AtlasSource(**self._get_source_kwargs())],
             is_oneoff=self._is_oneoff

ripe/atlas/tools/commands/measurement.py

@@ -22,7 +22,8 @@ def add_arguments(self):
     def run(self):
 
         try:
-            measurement = Measurement(id=self.arguments.id)
+            measurement = Measurement(
+                id=self.arguments.id, user_agent=self.user_agent)
         except APIResponseError:
             raise RipeAtlasToolsException(
                 "That measurement does not appear to exist")

ripe/atlas/tools/commands/measurements.py

@@ -115,7 +115,8 @@ def run(self):
             self.arguments.field = ("id", "type", "description", "status")
 
         filters = self._get_filters()
-        measurements = MeasurementRequest(return_objects=True, **filters)
+        measurements = MeasurementRequest(
+            return_objects=True, user_agent=self.user_agent, **filters)
         truncated_measurements = itertools.islice(
             measurements, self.arguments.limit)
 

ripe/atlas/tools/commands/probe.py

@@ -20,7 +20,7 @@ def add_arguments(self):
     def run(self):
 
         try:
-            probe = Probe(id=self.arguments.id)
+            probe = Probe(id=self.arguments.id, user_agent=self.user_agent)
         except APIResponseError:
             raise RipeAtlasToolsException(
                 "That probe does not appear to exist")

ripe/atlas/tools/commands/probes.py

@@ -173,7 +173,8 @@ def run(self):
             ))
 
         self.set_aggregators()
-        probes = ProbeRequest(return_objects=True, **filters)
+        probes = ProbeRequest(
+            return_objects=True, user_agent=self.user_agent, **filters)
         if self.arguments.limit:
             truncated_probes = itertools.islice(probes, self.arguments.limit)
         else:

ripe/atlas/tools/commands/report.py

@@ -84,7 +84,10 @@ def add_arguments(self):
 
     def _get_request(self):
 
-        kwargs = {"msm_id": self.arguments.measurement_id}
+        kwargs = {
+            "msm_id": self.arguments.measurement_id,
+            "user_agent": self.user_agent
+        }
         if self.arguments.probes:
             kwargs["probe_ids"] = self.arguments.probes
         if self.arguments.start_time:
@@ -99,7 +102,8 @@ def _get_request(self):
     def run(self):
 
         try:
-            measurement = Measurement(id=self.arguments.measurement_id)
+            measurement = Measurement(
+                id=self.arguments.measurement_id, user_agent=self.user_agent)
         except APIResponseError:
             raise RipeAtlasToolsException("That measurement does not exist")
 

ripe/atlas/tools/commands/stream.py

@@ -38,7 +38,8 @@ def add_arguments(self):
     def run(self):
 
         try:
-            measurement = Measurement(id=self.arguments.measurement_id)
+            measurement = Measurement(
+                id=self.arguments.measurement_id, user_agent=self.user_agent)
         except APIResponseError:
             raise RipeAtlasToolsException("That measurement does not exist")
 

tests/commands/report.py

@@ -1,3 +1,5 @@
+# coding=utf-8
+
 import mock
 import unittest
 
@@ -6,6 +8,7 @@
 from ripe.atlas.tools.commands.report import Command
 from ripe.atlas.tools.exceptions import RipeAtlasToolsException
 from ripe.atlas.tools.renderers import Renderer
+from ripe.atlas.tools.version import __version__
 from ..base import capture_sys_output
 
 
@@ -284,3 +287,20 @@ def test_asns_filter(self):
                     self.cmd.init_args(["1", "--probe-asns", "3334"])
                     self.cmd.run()
                     self.assertEquals(stdout.getvalue(), expected_output)
+
+    def test_user_agent(self):
+        standard = "RIPE Atlas Tools (Magellan) {}".format(__version__)
+        tests = {
+            standard: standard,
+            "Some custom agent": "Some custom agent",
+            "Some custom agent\nwith a second line": "Some custom agent",
+            "x" * 3000: "x" * 128,
+            "Πράκτορας χρήστη": "Πράκτορας χρήστη",
+            "이것은 테스트 요원": "이것은 테스트 요원",
+        }
+        self.assertEqual(self.cmd.user_agent, standard)
+        for in_string, out_string in tests.items():
+            path = "ripe.atlas.tools.commands.base.open"
+            content = mock.mock_open(read_data=in_string)
+            with mock.patch(path, content):
+                self.assertEqual(Command().user_agent, out_string)