tracker/services/result-queue/result_queue.py at update-cloudbuild · LearningAzure2020/tracker · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import os
import sys
import time
import requests
import logging
import json
import emoji
import traceback
import asyncio
from threading import Thread
from flask import Flask, request
from redis import Redis, ConnectionPool
from rq import Queue, Retry, Worker

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
pool = ConnectionPool(host="127.0.0.1", port=6379, db=0)

PROCESSOR_URL = "http://result-processor.scanners.svc.cluster.local"

redis = Redis(connection_pool=pool)
https_queue = Queue("https", connection=redis)
ssl_queue = Queue("ssl", connection=redis)
dns_queue = Queue("dns", connection=redis)

default_queues = {"https": https_queue, "ssl": ssl_queue, "dns": dns_queue}


def Server(process_name, queues=default_queues):

    flask_app = Flask(process_name)
    flask_app.config["queues"] = queues

    @flask_app.route("/https", methods=["POST"])
    def enqueue_https():
        logging.info("HTTPS result processing request received.")
        try:
            payload = request.get_json(force=True)
            designated_queue = flask_app.config["queues"].get("https", None)
            designated_queue.enqueue(
                dispatch_https,
                payload,
                retry=Retry(max=3),
                job_timeout=86400,
                result_ttl=86400,
            )
            msg = "HTTPS result processing request enqueued."
            logging.info(msg)
        except Exception as e:
            msg = f"An unexpected error occurred while attempting to enqueue HTTPS result processing request: ({type(e).__name__}: {str(e)})"
            logging.error(msg)
            logging.error(f"Full traceback: {traceback.format_exc()}")
        return msg

    @flask_app.route("/ssl", methods=["POST"])
    def enqueue_ssl():
        logging.info("SSL result processing request received.")
        try:
            payload = request.get_json(force=True)
            designated_queue = flask_app.config["queues"].get("ssl", None)
            designated_queue.enqueue(
                dispatch_ssl,
                payload,
                retry=Retry(max=3),
                job_timeout=86400,
                result_ttl=86400,
            )
            msg = "SSL result processing request enqueued."
            logging.info(msg)
        except Exception as e:
            msg = f"An unexpected error occurred while attempting to enqueue SSL result processing request: ({type(e).__name__}: {str(e)})"
            logging.error(msg)
            logging.error(f"Full traceback: {traceback.format_exc()}")
        return msg

    @flask_app.route("/dns", methods=["POST"])
    def enqueue_dns():
        logging.info("DNS result processing request received.")
        try:
            payload = request.get_json(force=True)
            designated_queue = flask_app.config["queues"].get("dns", None)
            designated_queue.enqueue(
                dispatch_dns,
                payload,
                retry=Retry(max=3),
                job_timeout=86400,
                result_ttl=86400,
            )
            msg = "DNS result processing request enqueued."
            logging.info(msg)
        except Exception as e:
            msg = f"An unexpected error occurred while attempting to enqueue DNS result processing request: ({type(e).__name__}: {str(e)})"
            logging.error(msg)
            logging.error(f"Full traceback: {traceback.format_exc()}")
        return msg

    return flask_app


app = Server(__name__)


def dispatch_https(payload):
    logging.info("Dispatching HTTPS result processing request")
    try:
        requests.post(PROCESSOR_URL, json=payload)
        return "Dispatched HTTPS result processing request."
    except Exception as e:
        msg = f"An unexpected error occurred while attempting to dispatch HTTPS result processing request: ({type(e).__name__}: {str(e)})"
        logging.error(msg)
        logging.error(f"Full traceback: {traceback.format_exc()}")
        return msg


def dispatch_ssl(payload):
    logging.info("Dispatching SSL result processing request")
    try:
        requests.post(PROCESSOR_URL, json=payload)
        return "Dispatched SSL result processing request."
    except Exception as e:
        msg = f"An unexpected error occurred while attempting to dispatch SSL result processing request: ({type(e).__name__}: {str(e)})"
        logging.error(msg)
        logging.error(f"Full traceback: {traceback.format_exc()}")
        return msg


def dispatch_dns(payload):
    logging.info("Dispatching DNS result processing request")
    try:
        requests.post(PROCESSOR_URL, json=payload)
        return "Dispatched DNS result processing request."
    except Exception as e:
        msg = f"An unexpected error occurred while attempting to dispatch DNS result processing request: ({type(e).__name__}: {str(e)})"
        logging.error(msg)
        logging.error(f"Full traceback: {traceback.format_exc()}")
        return msg