Completed
Push — master ( e72060...c61f62 )
by Ryan
01:01
created

traceroutedb.lookup_trace()   A

Complexity

Conditions 2

Size

Total Lines 8

Duplication

Lines 0
Ratio 0 %
Metric Value
cc 2
dl 0
loc 8
rs 9.4285
1
#!/usr/bin/env python
2
3
from __future__ import print_function
4
from flask import Flask, request, abort
5
from werkzeug.contrib.cache import SimpleCache
6
import logging
7
import sys
8
import json
9
import geoip2.database
10
11
try:
12
    import psycopg2
13
    from psycopg2 import extras as pgextras
14
except ImportError:
15
    print("psycopg2 needed to run server")
16
    sys.exit(1)
17
18
app = Flask(__name__)
19
cache = SimpleCache()
20
logger = logging.getLogger(__name__)
21
22
23
CACHE_TIMEOUT = 300
24
25
26
class cached(object):
27
28
    def __init__(self, timeout=None):
29
        self.timeout = timeout or CACHE_TIMEOUT
30
31
    def __call__(self, f):
32
        def decorator(*args, **kwargs):
33
            response = cache.get(request.path)
34
            if response is None:
35
                response = f(*args, **kwargs)
36
                cache.set(request.path, response, self.timeout)
37
            return response
38
        return decorator
39
40
41
def dictToHstore(in_dict):
42
    single_kvs = []
43
    for key in in_dict.iterkeys():
44
        if in_dict[key]:
45
            single_kvs.append("{0}=>{1}".format(key, in_dict[key]))
46
    hstore_str = "'{0}'".format(", ".join(single_kvs))
47
    return hstore_str
48
49
50
@app.route("/rules", methods=["GET"])
51
@cached()
52
def get_rules():
53
    try:
54
        connv = psycopg2.connect("dbname=traceroutedb user=postgres host=localhost")
55
    except psycopg2.OperationalError as e:
56
        logging.error(str(e))
57
        abort(503)
58
    cur = connv.cursor()
59
    cur.execute("SELECT * FROM endpoints;")
60
    rows = cur.fetchall()
61
    connv.close()
62
    return json.dumps(rows)
63
64
65
@app.route("/trace", methods=["POST"])
66
def receive_traces():
67
    config = app.config["trdb"]
68
    if config.get("mmdb", False):
69
        reader = config["mmdb"]
70
71
    if request.method == "POST":
72
        cur = conn.cursor()
73
        data = request.get_json(force=True)
74
75
        trace = data["data"]
76
        reporter = data["reporter"]
77
        kvs = {"note": data.get("note"), "ext_ip": data.get("ext_ip")}
78
79
        if config.debug:
80
            print("SELECT nextval('traceroute_id_seq');")
81
            trace_id = "DEBUG_ID"
82
        else:
83
            try:
84
                cur.execute("SELECT nextval('traceroute_id_seq');")
85
            except psycopg2.ProgrammingError as e:
86
                logging.error(str(e))
87
                conn.rollback()
88
                abort(503)
89
            trace_id = cur.fetchone()[0]
90
91
        trace_sql = "INSERT INTO traceroute VALUES ({0}, '{1}', '{2}', now(), '{3}', {4}::HSTORE);".format(trace_id, trace["src_ip"], trace["dst_ip"], reporter, dictToHstore(kvs))
92
        if config.debug:
93
            print(trace_sql)
94
        else:
95
            try:
96
                cur.execute(trace_sql)
97
            except psycopg2.ProgrammingError as e:
98
                logging.error(str(e))
99
                conn.rollback()
100
                abort(503)
101
102
        hops = trace["hops"]
103
        for key in hops.keys():
104
            hop = hops[key]
105
            for probe in hop:
106
                if probe["ip"] is None or probe["rtt"] is None:
107
                    continue
108
                else:
109
                    try:
110
                        probe["isp"] = reader.isp(probe["ip"])
111
                    except:
112
                        pass
0 ignored issues
show
Unused Code introduced by
This except handler seems to be unused and could be removed.

Except handlers which only contain pass and do not have an else clause can usually simply be removed:

try:
    raises_exception()
except:  # Could be removed
    pass
Loading history...
113
114
                kvs = {}
115
                time = probe.get("rtt", None)
116
                if time is not None and time != "None":
117
                    kvs["time"] = time
118
119
                anno = probe.get("anno", None)
120
                if anno:
121
                    kvs["anno"] = anno
122
123
                try:
124
                    asn = probe.get("isp").raw.get("autonomous_system_number", None)
125
                    if asn:
126
                        kvs["asn"] = asn
127
                except:
128
                    pass
0 ignored issues
show
Unused Code introduced by
This except handler seems to be unused and could be removed.

Except handlers which only contain pass and do not have an else clause can usually simply be removed:

try:
    raises_exception()
except:  # Could be removed
    pass
Loading history...
129
130
                kvs = dictToHstore(kvs)
131
                if config.debug:
132
                    print("INSERT INTO hop VALUES (nextval('probe_id_seq'), {0}, {1}, {2}, '{3}', now());".format(trace_id, key, kvs, probe["ip"]))
133
                else:
134
                    try:
135
                        cur.execute("INSERT INTO hop VALUES (nextval('probe_id_seq'), {0}, {1}, {2}, '{3}', now());".format(trace_id, key, kvs, probe["ip"]))
136
                    except psycopg2.ProgrammingError as e:
137
                        logging.error(str(e))
138
                        conn.rollback()
139
                        abort(503)
140
    conn.commit()
141
    cur.close()
142
    conn.commit()
143
    print("ip:", request.remote_addr, "submitted result for:", trace["src_ip"], ">", trace["dst_ip"], "trace_id:", trace_id)
144
    return 'OK'
145
146
147
@app.route("/trace/<trace_id>", methods=["GET"])
148
def lookup_trace(trace_id):
149
    cur = conn.cursor(cursor_factory=pgextras.DictCursor)
150
    pgextras.register_hstore(cur)
151
    cur.execute("SELECT * FROM trv_trace WHERE traceroute_id = {id} ORDER BY traceroute_id,hop_number;".format(id=trace_id))
152
    rows = cur.fetchall()
153
    out = [json.dumps(x) for x in rows]
154
    return "\n".join(out)
155
156
157
try:
158
    conn = psycopg2.connect("dbname=traceroutedb user=postgres host=localhost")
159
except psycopg2.OperationalError as e:
160
    logging.error(str(e))
161
    sys.exit(1)
162
163
164
def run_server(config, app=app):
0 ignored issues
show
Comprehensibility Bug introduced by
app is re-defining a name which is already available in the outer-scope (previously defined on line 18).

It is generally a bad practice to shadow variables from the outer-scope. In most cases, this is done unintentionally and might lead to unexpected behavior:

param = 5

class Foo:
    def __init__(self, param):   # "param" would be flagged here
        self.param = param
Loading history...
165
    if config.debug:
166
        logger.setLevel(logging.DEBUG)
167
168
    app.config["trdb"] = config
169
    if config.mmdb:
170
        reader = geoip2.database.Reader(config.mmdb)
171
        app.config["trdb"]["mmdb"] = reader
172
173
    app.run(port=9001, debug=config.debug, host='::')
174