ail-framework/doc/api/create_lu_graph.py

#!/usr/bin/env python3
# -*-coding:UTF-8 -*

'''
lu
'''

import os
import sys
import redis
import datetime
import heapq
import operator
import matplotlib.pyplot as plt
import numpy as np

sys.path.append(os.environ['AIL_BIN'])

from Helper import Process

def create_pie_chart(db_key, date, pie_title, path, save_name):

    monthly_credential_by_tld = server_statistics.hkeys(db_key + date)

    l_tld = []
    for tld in monthly_credential_by_tld:
        nb_tld = server_statistics.hget(db_key + date, tld)
        if nb_tld is not None:
            nb_tld = int(nb_tld)
        else:
            nb_tld = 0
        l_tld.append( (tld, nb_tld) )

    mail_tld_top5 = heapq.nlargest(5, l_tld, key=operator.itemgetter(1))

    # Pie chart, where the slices will be ordered and plotted counter-clockwise:
    labels = []
    sizes = []
    explode = []  # only "explode" the 2nd slice (i.e. 'Hogs')
    for tld in mail_tld_top5:
        labels.append(tld[0])
        sizes.append(tld[1])
        explode.append(0)

    nb_tld = server_statistics.hget(db_key + date, 'lu')
    if nb_tld is not None:
        nb_tld = int(nb_tld)
    else:
        nb_tld = 0
    labels.append('lu')
    sizes.append(nb_tld)
    explode.append(0.3)  # only "explode" lu slice
    explode = tuple(explode)

    fig1, ax1 = plt.subplots()
    ax1.pie(sizes, explode=explode, labels=labels, autopct='%1.1f%%',
            shadow=True, startangle=90)
    ax1.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle.

    ax1.set_title(pie_title)
    #plt.show()
    plt.savefig(os.path.join(path, save_name))
    plt.close(fig1)

def create_donut_chart(db_key, date, pie_title, path, save_name):

    monthly_credential_by_tld = server_statistics.hkeys(db_key + date)
    print()

    l_tld = []
    for tld in monthly_credential_by_tld:
        nb_tld = server_statistics.hget(db_key + date, tld)
        if nb_tld is not None:
            nb_tld = int(nb_tld)
        else:
            nb_tld = 0
        l_tld.append( (tld, nb_tld) )

    mail_tld_top5 = heapq.nlargest(5, l_tld, key=operator.itemgetter(1))

    # Pie chart, where the slices will be ordered and plotted counter-clockwise:
    recipe = []
    data = []
    for tld in mail_tld_top5:
        recipe.append(tld[0])
        data.append(tld[1])

    nb_tld = server_statistics.hget(db_key + date, 'lu')
    if nb_tld is not None:
        nb_tld = int(nb_tld)
    else:
        nb_tld = 0
    recipe.append('lu')
    data.append(nb_tld)

    fig1, ax1 = plt.subplots(figsize=(6, 3), subplot_kw=dict(aspect="equal"))

    wedges, texts = ax1.pie(data, wedgeprops=dict(width=0.5), startangle=-40)

    bbox_props = dict(boxstyle="square,pad=0.3", fc="w", ec="k", lw=0.72)
    kw = dict(xycoords='data', textcoords='data', arrowprops=dict(arrowstyle="-"),
              bbox=bbox_props, zorder=0, va="center")

    for i, p in enumerate(wedges):
        ang = (p.theta2 - p.theta1)/2. + p.theta1
        y = np.sin(np.deg2rad(ang))
        x = np.cos(np.deg2rad(ang))
        horizontalalignment = {-1: "right", 1: "left"}[int(np.sign(x))]
        connectionstyle = "angle,angleA=0,angleB={}".format(ang)
        kw["arrowprops"].update({"connectionstyle": connectionstyle})
        ax1.annotate(recipe[i], xy=(x, y), xytext=(1.35*np.sign(x), 1.4*y),
                     horizontalalignment=horizontalalignment, **kw)

    ax1.set_title(pie_title)
    #plt.show()
    plt.savefig(os.path.join(path, save_name))
    plt.close(fig1)

if __name__ == '__main__':

    path = os.path.join(os.environ['AIL_HOME'], 'doc') # path to module config file

    config_section = 'ARDB_Statistics'

    p = Process(config_section, False)

    # ARDB #
    server_statistics = redis.StrictRedis(
        host=p.config.get("ARDB_Statistics", "host"),
        port=p.config.getint("ARDB_Statistics", "port"),
        db=p.config.getint("ARDB_Statistics", "db"),
        decode_responses=True)

    date = datetime.datetime.now().strftime("%Y%m")
    create_pie_chart('credential_by_tld:', date, "AIL: Credential leak by tld", path, 'AIL_credential_by_tld.png')
    create_pie_chart('mail_by_tld:', date, "AIL: mail leak by tld", path, 'AIL_mail_by_tld.png')
    create_pie_chart('SQLInjection_by_tld:', date, "AIL: sqlInjection by tld", path, 'AIL_sqlInjectionl_by_tld.png')