#!/usr/bin/env python3
import os
import sys
from dateutil import parser
import json, logging
from flask import Flask, request
from rocketchat_API.rocketchat import RocketChat
from requests import sessions

listenPort = 29120
listenIP = '::1'
botUser = '{{ prometheus_alertmanager.rocketchatbot.user }}'
botPass = '{{ prometheus_alertmanager.rocketchatbot.pass }}'
botChatURL = '{{ prometheus_alertmanager.rocketchatbot.url }}'

# ignore fields while printing labels
label_ignore_fields = ['ansible_group_ungrouped']

icons = {
    'firing': ':warning:',
    'resolved': ':white_check_mark:'
}

'''
Example attachment:

message.attachment = [{
    "color": "#ff0000",
    "text": "Yay for gruggy!",
    "ts": "2016-12-09T16:53:06.761Z",
    "thumb_url": "data:image/gif;base64,...",  # add thumb image as base64 if needed
    "message_link": "https://google.com",
    "collapsed": false,
    "author_name": "Bradley Hilton",
    "author_link": "https://rocket.chat/",
    "author_icon": "https://avatars.githubusercontent.com/u/850391?v=3",
    "title": "Attachment Example",
    "title_link": "https://youtube.com",
    "title_link_download": true,
    "image_url": "http://res.guggy.com/logo_128.png",
    "audio_url": "http://www.w3schools.com/tags/horse.mp3",
    "video_url": "http://www.w3schools.com/tags/movie.mp4",
    "fields": [{
      "short": true,
      "title": "Test",
      "value": "Testing out something or other"
    },{
      "short": true,
      "title": "Another Test",
      "value": "[Link](https://google.com/) something and this and that."
    }]
}]


Example alert json:

[{
    "receiver": "default", 
    "status": "resolved", 
    "alerts": [{
        "status": "resolved", 
        "labels": {
            "alertname": "ProbeFailed", 
            "ansible_group_ungrouped": "1", 
            "instance": "monitoring.cccv.de", 
            "job": "blackbox", 
            "module": "http_2xx", 
            "scraper": "monitoring.cccv.de", 
            "severity": "critical", 
            "target": "https://shells.darmstadt.ccc.de/~psy/alerttest.html"
        }, 
        "annotations": {
            "description": "monitoring.cccv.de probe for https://shells.darmstadt.ccc.de/~psy/alerttest.html failed", 
            "title": "monitoring.cccv.de: target https://shells.darmstadt.ccc.de/~psy/alerttest.html failed"
        }, 
        "startsAt": "2021-01-04T15:36:30.85487796Z", 
        "endsAt": "2021-01-04T17:18:30.85487796Z", 
        "generatorURL": "http://monitoring:29090/graph?g0.expr=probe_success%7Binstance%3D%22monitoring.cccv.de%22%2Cjob%3D%22blackbox%22%7D+%3D%3D+0+or+absent%28probe_success%7Binstance%3D%22monitoring.cccv.de%22%2Cjob%3D%22blackbox%22%7D%29&g0.tab=1"
    }], 
    "groupLabels": {
        "alertname": "ProbeFailed"
    }, 
    "commonLabels": {
        "alertname": "ProbeFailed", 
        "ansible_group_ungrouped": "1", 
        "instance": "monitoring.cccv.de", 
        "job": "blackbox", 
        "module": "http_2xx", 
        "scraper": "monitoring.cccv.de", 
        "severity": "critical", 
        "target": "https://shells.darmstadt.ccc.de/~psy/alerttest.html"
    }, 
    "commonAnnotations": {
        "description": "monitoring.cccv.de probe for https://shells.darmstadt.ccc.de/~psy/alerttest.html failed", 
        "title": "monitoring.cccv.de: target https://shells.darmstadt.ccc.de/~psy/alerttest.html failed"
    }, 
    "externalURL": "http://monitoring:29093", 
    "version": "4", 
    "groupKey": "{}:{alertname=\"ProbeFailed\"}"
}]
'''

app = Flask(__name__)
app.secret_key = os.urandom(128)

api_session = sessions.Session()
api = RocketChat(botUser, botPass, server_url=botChatURL, session=api_session)


@app.route('/<chatName>/alert', methods=['POST'])
def postAlertmanager(chatName):
    try:
        content = json.loads(request.get_data())
    except Exception as error:
        api.chat_post_message("Error to read json: " + str(error), channel=chatName, alias='Alertmanager')
        app.logger.info("\t%s", error)
        return "Alert fail", 200

    success = True
    for alert in content['alerts']:
        try:
            send_alert_message(chatName, alert)
        except Exception as error:
            success = False
            app.logger.info("\t%s", error)
            api.chat_post_message("Error to post alert: " + str(error), channel=chatName, alias='Alertmanager')

    if success:
        return "Alert OK", 200
    else:
        return "Alert fail", 200


def send_alert_message(channel, alert):
    message = "[**{state}**] [{instance}] {amendment}".format(state=alert['status'],
                                                         instance=alert['labels'].get('instance', 'unknown'),
                                                         amendment='@here' if alert['status'] == 'firing' else alert['annotations'].get('title'))
    attach = []

    # attach details if alert is firing
    if alert['status'] == 'firing':
        attach.append({
            "color": "#ff0000",
            "text": "\n".join(["**{}**: {}".format(key, value) for key, value in alert['annotations'].items() if
                               key not in ['title']]),
            "ts": parser.parse(alert['startsAt'] if alert['status'] == 'firing' else alert['endsAt']).isoformat(),
            "message_link": alert.get('generatorURL', '#'),  # link for timestamp, mandatory to display timestamp
            "collapsed": True,  # collapse details by default
            "author_name": alert['labels'].get('instance', 'unknown'),
            "title": alert['annotations']['title'],
            "fields": [{"short": True, "title": key, "value": value} for key, value in alert['labels'].items() if
                       key not in label_ignore_fields]
        })

    res = api.chat_post_message(message, attachments=attach, channel=channel,
                                alias=alert['labels'].get('alertname', 'unknown alert'),
                                emoji=icons.get(alert['status']))

    if not res.ok:
        raise Exception(res.content)


if __name__ == '__main__':
    if len(sys.argv) == 1:
        logging.basicConfig(level=logging.INFO)
        app.run(host=listenIP, port=listenPort)
    else:
        from gevent.pywsgi import WSGIServer

        http_server = WSGIServer((listenIP, listenPort), app)
        http_server.serve_forever()