Auxiliary-Injector/listener.py at main · SSantosLab/Auxiliary-Injector · 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
"""
Python script to listen from non-LVK alerts trough gcn.
"""
from tests.tests import parse_json_from_bytes
import glob
import time
import traceback
import json
from argparse import ArgumentParser
import yaml
from yaml.loader import SafeLoader
from gcn_kafka import Consumer
from handlers.streamer import Alert
from handlers.emails import EmailBot
from handlers.slack import SlackBot
import datetime
from tests.tests import testHandler
from handlers.streamer import handle

allowedModes = ["test","all","test g","test x","test v"]
allowedChars = ["g","x","v"]

# Gamma ray: Einstein probe  gcn.notices.einstein_probe.wxt.alert
# X ray: Swift gcn.notices.swift.bat.guano
# Radio (r): No great option, yet
# Neutrino: IceCube : gcn.notices.icecube.lvk_nu_track_search

subscriptionDict = {
                    "g":"gcn.notices.einstein_probe.wxt.alert",
                    "x":"gcn.notices.swift.bat.guano",
                    "v":"gcn.notices.icecube.lvk_nu_track_search"
                    }

# Other sources to consider: SuperNova Early Warning System (SNEWS)

def elapsedTimeString(start):
    elapsed = int(time.time() - start)
    return "{}h {:02d}m {:02d}s".format(elapsed//(60*60), elapsed//60%60, elapsed%60)

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument('--mode',
                        default='all',
                        help='Chose a `mode` for listen to gcn alerts. ' +\
                        'Default mode is `all`, which will include gamma ray (g), x-ray (x), and neutrino events (v).' +\
                        'For offline testing of all modes, use `test` mode. ' +\
                        'For offline testing of one mode, use `test P`, where `P` is one of the chars associated with a different messenger',
                        choices=allowedModes)

    args = parser.parse_args()
    mode = args.mode
    start_time = time.time()
    print("Listener Activating")
    print('Running Listener in {} mode...'.format(mode), flush=True)

    # alert_streamer = alertStreamer(mode=mode)
    email_bot = EmailBot(mode=mode)
    slack_bot = SlackBot(mode=mode)

    with open('configs/gcn_credentials.yaml', 'r', encoding='utf-8') as f:
        gcn = yaml.load(f, Loader=SafeLoader)
        slack_bot.post_message("","Starting `listener.py` for auxiliary injector")

    if mode == 'test':
        print('Reading test events...')
        for character in allowedChars:
            postString = "Starting test handler for mode {}".format(character)
            slack_bot.post_message("",postString)
            print(postString)
            testHandler(character)

    elif mode.startswith('test'):
        testModes = list(mode.split(" ")[1]) # These are the modes that are to be started for testing, all chars

        for mode in testModes:
            postString = "Starting test handler for mode {}".format(mode)
            slack_bot.post_message("",postString)
            print(postString)
            testHandler(mode)


    elif mode=='all':
        try:
            consumer = Consumer(client_id=gcn['client_id'],
                                client_secret=gcn['client_secret'])
            # initial
            for value in subscriptionDict.values():
                consumer.subscribe([value])

            today = datetime.date.today().day
            init_day = 1

            while True:
                for message in consumer.consume(timeout=10):
                    print('Trigger Received...')
                    if message.value()['type']=='IceCube LVK Alert Nu Track Search': # For the one strangely formatted gcn type
                        gcn_alert=parse_json_from_bytes(message.value())
                    else:
                        gcn_alert = json.loads(message.value())
#                    gcn_alert = dict(message.value())
                    print('Passing event to Handler.', flush=True)
                    print(gcn_alert)
                    handle(gcn_alert,"placeholder")

                if datetime.date.today().day != today:
                    slack_bot.post_message("","`listener.py` for `Auxiliary-Injector` has been running nonstop in {} mode for {} days".format(mode, init_day))
                    today = datetime.date.today().day
                    init_day +=1

        except Exception as e:
            print(e)
            slack_bot.post_message("","Listener went down! Please investigate! Traceback attached <@U05V24X6MHB><@U0545QECWJZ><@UAV5VNB9N>. {}".format(traceback.format_exc()))
            email_bot.send_email(subject='Listener went down, see traceback',
                                text=traceback.format_exc(),
                                emergency=True)
    else:
        raise ValueError
        print("Mode not supported.\nSupplied mode: {}.\nAllowed modes are {}.".format(mode,allowedModes))