diff --git a/ripe/atlas/tools/commands/base.py b/ripe/atlas/tools/commands/base.py index 43c2cc8..40adae4 100644 --- a/ripe/atlas/tools/commands/base.py +++ b/ripe/atlas/tools/commands/base.py @@ -1,4 +1,5 @@ import argparse +import os import re import six import sys @@ -6,6 +7,7 @@ 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.read().strip() + except IOError: + pass # We go with the default + + return "RIPE Atlas Tools (Magellan) {}".format(__version__) + class TabularFieldsMixin(object): """ diff --git a/ripe/atlas/tools/commands/measure/base.py b/ripe/atlas/tools/commands/measure/base.py index 989dc26..b1b3c16 100644 --- a/ripe/atlas/tools/commands/measure/base.py +++ b/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 diff --git a/ripe/atlas/tools/commands/measurement.py b/ripe/atlas/tools/commands/measurement.py index 2404a5c..8233f51 100644 --- a/ripe/atlas/tools/commands/measurement.py +++ b/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") diff --git a/ripe/atlas/tools/commands/measurements.py b/ripe/atlas/tools/commands/measurements.py index 299eb1c..537a41c 100644 --- a/ripe/atlas/tools/commands/measurements.py +++ b/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) diff --git a/ripe/atlas/tools/commands/probe.py b/ripe/atlas/tools/commands/probe.py index eab6be6..764fcf6 100644 --- a/ripe/atlas/tools/commands/probe.py +++ b/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") diff --git a/ripe/atlas/tools/commands/probes.py b/ripe/atlas/tools/commands/probes.py index 0286fee..4f0c40c 100644 --- a/ripe/atlas/tools/commands/probes.py +++ b/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: diff --git a/ripe/atlas/tools/commands/report.py b/ripe/atlas/tools/commands/report.py index 4ff716c..d28c9da 100644 --- a/ripe/atlas/tools/commands/report.py +++ b/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") diff --git a/ripe/atlas/tools/commands/stream.py b/ripe/atlas/tools/commands/stream.py index aef1b8f..bb27dc1 100644 --- a/ripe/atlas/tools/commands/stream.py +++ b/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") diff --git a/tests/commands/report.py b/tests/commands/report.py index 6e2c09f..3926abb 100644 --- a/tests/commands/report.py +++ b/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,18 @@ 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, + "Some custom agent", + "Αυτό είναι ένας παράγοντας δοκιμή", + "이것은 테스트 요원", + ] + self.assertEqual(self.cmd.user_agent, standard) + for agent in tests: + path = "ripe.atlas.tools.commands.base.open" + content = mock.mock_open(read_data=agent) + with mock.patch(path, content): + self.assertEqual(Command().user_agent, agent)