Untitled

import psutil
import json
import requests

INPUT_JSON = input()
SHUFFLE_WEBHOOK_URL = "http://192.168.0.106:3001/api/v1/hooks/webhook_1fc7bc57-4cd1-4cf7-9e42-126ac71f724b"

def get_memory_info():
    return {
        "total_memory": psutil.virtual_memory().total / (1024.0 ** 3),
        "available_memory": psutil.virtual_memory().available / (1024.0 ** 3),
        "used_memory": psutil.virtual_memory().used / (1024.0 ** 3),
        "memory_percentage": psutil.virtual_memory().percent
    }

def get_cpu_info():
    return {
        "physical_cores": psutil.cpu_count(logical=False),
        "total_cores": psutil.cpu_count(logical=True),
        "processor_speed": psutil.cpu_freq().current,
        "cpu_usage_per_core": dict(enumerate(psutil.cpu_percent(percpu=True, interval=1))),
        "total_cpu_usage": psutil.cpu_percent(interval=1)
    }

def get_disk_info():
    partition_usage = psutil.disk_usage('/')
    return {
       "total_space": partition_usage.total / (1024.0 ** 3),
       "used_space": partition_usage.used / (1024.0 ** 3),
       "free_space": partition_usage.free / (1024.0 ** 3),
       "usage_percentage": partition_usage.percent
    }

def get_alert_info():
    disk_data = json.loads(INPUT_JSON)
    parameters = disk_data.get("parameters").get("extra_args")
    return {
        "text": parameters[0],
        "timestamp": parameters[1],
	"srcip": parameters[2],
	"dstip": parameters[3],
	"agent_id": parameters[4],
	"agent_name": parameters[5],
	"agent_ip": parameters[6],
	"api_key": parameters[7]
    }

def send_to_shuffle(data):
  headers = {'Content-Type': 'application/json'}
  payload = {'system': json.dumps(data, indent=4)}

  response = requests.post(SHUFFLE_WEBHOOK_URL, headers=headers, json=payload)

if __name__ == '__main__':
  under_attack = False
  cpu_info = get_cpu_info()
  memory_info = get_memory_info()
  
  if (cpu_info["total_cpu_usage"] > 90):
    under_attack = True
  elif (memory_info["memory_percentage"] > 90):
    under_attack = True

  data = {
    "cpu_info": cpu_info,
    "memory_info": memory_info,
    "disk_info": get_disk_info(),
    "alert_info": get_alert_info(),
    "under_attack": under_attack
  }

  send_to_shuffle(data)