1# Copyright (c) 2015 RIPE NCC
2#
3# This program is free software: you can redistribute it and/or modify
4# it under the terms of the GNU General Public License as published by
5# the Free Software Foundation, either version 3 of the License, or
6# (at your option) any later version.
7#
8# This program is distributed in the hope that it will be useful,
9# but WITHOUT ANY WARRANTY; without even the implied warranty of
10# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11# GNU General Public License for more details.
12#
13# You should have received a copy of the GNU General Public License
14# along with this program.  If not, see <http://www.gnu.org/licenses/>.
15from datetime import datetime
16from dateutil.tz import tzutc
17
18from .request import AtlasRequest
19from .exceptions import CousteauGenericError, APIResponseError
20
21
22class EntityRepresentation(object):
23    """
24    A crude representation of entity's meta data as we get it from the API.
25    """
26
27    API_META_URL = ""
28
29    def __init__(self, **kwargs):
30
31        self.id = kwargs.get("id")
32        self.server = kwargs.get("server")
33        self.verify = kwargs.get("verify", True)
34        self.api_key = kwargs.get("key", "")
35        self.meta_data = kwargs.get("meta_data")
36        self._user_agent = kwargs.get("user_agent")
37        self._fields = kwargs.get("fields")
38        self.get_params = {}
39
40        if self.meta_data is None and self.id is None:
41            raise CousteauGenericError(
42                "Id or meta_data should be passed in order to create object."
43            )
44
45        if self._fields:
46            self.update_get_params()
47
48        if self.meta_data is None:
49            if not self._fetch_meta_data():
50                raise APIResponseError(self.meta_data)
51
52        self._populate_data()
53
54    def update_get_params(self):
55        """Update HTTP GET params with the given fields that user wants to fetch."""
56        if isinstance(self._fields, (tuple, list)):  # tuples & lists > x,y,z
57            self.get_params["fields"] = ",".join([str(_) for _ in self._fields])
58        elif isinstance(self._fields, str):
59            self.get_params["fields"] = self._fields
60
61    def _fetch_meta_data(self):
62        """Makes an API call to fetch meta data for the given probe and stores the raw data."""
63        is_success, meta_data = AtlasRequest(
64            url_path=self.API_META_URL.format(self.id),
65            key=self.api_key,
66            server=self.server,
67            verify=self.verify,
68            user_agent=self._user_agent
69        ).get(**self.get_params)
70
71        self.meta_data = meta_data
72        if not is_success:
73            return False
74
75        return True
76
77    def _populate_data(self):
78        """
79        Passing some raw meta data from API response to instance properties
80        """
81        raise NotImplementedError()
82
83
84class Probe(EntityRepresentation):
85    """
86    A crude representation of probe's meta data as we get it from the API.
87    """
88    API_META_URL = "/api/v2/probes/{0}/"
89
90    def _populate_data(self):
91        """Assing some probe's raw meta data from API response to instance properties"""
92        if self.id is None:
93            self.id = self.meta_data.get("id")
94        self.is_anchor = self.meta_data.get("is_anchor")
95        self.country_code = self.meta_data.get("country_code")
96        self.description = self.meta_data.get("description")
97        self.is_public = self.meta_data.get("is_public")
98        self.asn_v4 = self.meta_data.get("asn_v4")
99        self.asn_v6 = self.meta_data.get("asn_v6")
100        self.address_v4 = self.meta_data.get("address_v4")
101        self.address_v6 = self.meta_data.get("address_v6")
102        self.prefix_v4 = self.meta_data.get("prefix_v4")
103        self.prefix_v6 = self.meta_data.get("prefix_v6")
104        self.geometry = self.meta_data.get("geometry")
105        self.tags = self.meta_data.get("tags")
106        self.status = self.meta_data.get("status", {}).get("name")
107
108    def __str__(self):
109        return "Probe #{0}".format(self.id)
110
111    def __repr__(self):
112        return str(self)
113
114
115class Measurement(EntityRepresentation):
116    """
117    A crude representation of measurement's meta data as we get it from the API.
118    """
119    API_META_URL = "/api/v2/measurements/{0}/"
120
121    def _populate_data(self):
122        """Assinging some measurement's raw meta data from API response to instance properties"""
123        if self.id is None:
124            self.id = self.meta_data.get("id")
125
126        self.stop_time = None
127        self.creation_time = None
128        self.start_time = None
129        self.populate_times()
130        self.protocol = self.meta_data.get("af")
131        self.target_ip = self.meta_data.get("target_ip")
132        self.target_asn = self.meta_data.get("target_asn")
133        self.target = self.meta_data.get("target")
134        self.description = self.meta_data.get("description")
135        self.is_oneoff = self.meta_data.get("is_oneoff")
136        self.is_public = self.meta_data.get("is_public")
137        self.interval = self.meta_data.get("interval")
138        self.resolve_on_probe = self.meta_data.get("resolve_on_probe")
139        self.status_id = self.meta_data.get("status", {}).get("id")
140        self.status = self.meta_data.get("status", {}).get("name")
141        self.type = self.get_type()
142        self.result_url = self.meta_data.get("result")
143
144    def get_type(self):
145        """
146        Getting type of measurement keeping backwards compatibility for
147        v2 API output changes.
148        """
149        mtype = None
150        if "type" not in self.meta_data:
151            return mtype
152
153        mtype = self.meta_data["type"]
154        if isinstance(mtype, dict):
155            mtype = self.meta_data.get("type", {}).get("name", "").upper()
156        elif isinstance(mtype, str):
157            mtype = mtype
158
159        return mtype
160
161    def populate_times(self):
162        """
163        Populates all different meta data times that comes with measurement if
164        they are present.
165        """
166        stop_time = self.meta_data.get("stop_time")
167        if stop_time:
168            stop_naive = datetime.utcfromtimestamp(stop_time)
169            self.stop_time = stop_naive.replace(tzinfo=tzutc())
170
171        creation_time = self.meta_data.get("creation_time")
172        if creation_time:
173            creation_naive = datetime.utcfromtimestamp(creation_time)
174            self.creation_time = creation_naive.replace(tzinfo=tzutc())
175
176        start_time = self.meta_data.get("start_time")
177        if start_time:
178            start_naive = datetime.utcfromtimestamp(start_time)
179            self.start_time = start_naive.replace(tzinfo=tzutc())
180
181    def __str__(self):
182        return "Measurement #{0}".format(self.id)
183
184    def __repr__(self):
185        return str(self)
186