244 lines
8.3 KiB
Python
244 lines
8.3 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
# Copyright (c) 2021 Battelle Energy Alliance, LLC. All rights reserved.
|
|
|
|
import argparse
|
|
import ipaddress
|
|
import json
|
|
import os
|
|
import socket
|
|
import ssl
|
|
import subprocess
|
|
import sys
|
|
import urllib.request
|
|
|
|
from base64 import b64encode
|
|
from bs4 import BeautifulSoup
|
|
from bs4.element import Comment
|
|
from contextlib import closing
|
|
from http.client import HTTPSConnection, HTTPConnection
|
|
from multiprocessing import RawValue
|
|
from threading import Lock
|
|
|
|
NIC_BLINK_SECONDS = 10
|
|
|
|
###################################################################################################
|
|
class CancelledError(Exception):
|
|
"""Raised when user cancels the operation"""
|
|
pass
|
|
|
|
###################################################################################################
|
|
class Iface:
|
|
def __init__(self, name, description):
|
|
self.name = name
|
|
self.description = description
|
|
|
|
###################################################################################################
|
|
# clear the terminal window and exit the script
|
|
def clearquit():
|
|
os.system('clear')
|
|
sys.exit(0)
|
|
|
|
###################################################################################################
|
|
# print to stderr
|
|
def eprint(*args, **kwargs):
|
|
print(*args, file=sys.stderr, **kwargs)
|
|
|
|
###################################################################################################
|
|
# urlencode each character of a string
|
|
def aggressive_url_encode(string):
|
|
return "".join("%{0:0>2}".format(format(ord(char), "x")) for char in string)
|
|
|
|
###################################################################################################
|
|
# strip a prefix from the beginning of a string if needed
|
|
def remove_prefix(text, prefix):
|
|
if (len(prefix) > 0) and text.startswith(prefix):
|
|
return text[len(prefix):]
|
|
else:
|
|
return text
|
|
|
|
###################################################################################################
|
|
# nice human-readable file sizes
|
|
def sizeof_fmt(num, suffix='B'):
|
|
for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
|
|
if abs(num) < 1024.0:
|
|
return "%3.1f%s%s" % (num, unit, suffix)
|
|
num /= 1024.0
|
|
return "%.1f%s%s" % (num, 'Yi', suffix)
|
|
|
|
###################################################################################################
|
|
# convenient boolean argument parsing
|
|
def str2bool(v):
|
|
if v.lower() in ('yes', 'true', 't', 'y', '1'):
|
|
return True
|
|
elif v.lower() in ('no', 'false', 'f', 'n', '0'):
|
|
return False
|
|
else:
|
|
raise argparse.ArgumentTypeError('Boolean value expected.')
|
|
|
|
###################################################################################################
|
|
# will it float?
|
|
def isfloat(value):
|
|
try:
|
|
float(value)
|
|
return True
|
|
except ValueError:
|
|
return False
|
|
|
|
###################################################################################################
|
|
# check a string or list to see if something is a valid IP address
|
|
def isipaddress(value):
|
|
result = True
|
|
try:
|
|
if isinstance(value, list) or isinstance(value, tuple) or isinstance(value, set):
|
|
for v in value:
|
|
ip = ipaddress.ip_address(v)
|
|
else:
|
|
ip = ipaddress.ip_address(value)
|
|
except:
|
|
result = False
|
|
return result
|
|
|
|
###################################################################################################
|
|
# execute a shell process returning its exit code and output
|
|
def run_process(command, stdout=True, stderr=False, stdin=None, timeout=60):
|
|
retcode = -1
|
|
output = []
|
|
p = subprocess.run([command], input=stdin, universal_newlines=True, capture_output=True, shell=True, timeout=timeout)
|
|
if p:
|
|
retcode = p.returncode
|
|
if stderr and p.stderr:
|
|
output.extend(p.stderr.splitlines())
|
|
if stdout and p.stdout:
|
|
output.extend(p.stdout.splitlines())
|
|
|
|
return retcode, output
|
|
|
|
|
|
def tag_visible(element):
|
|
if element.parent.name in ['style', 'script', 'head', 'title', 'meta', '[document]']:
|
|
return False
|
|
if isinstance(element, Comment):
|
|
return False
|
|
return True
|
|
|
|
def text_from_html(body):
|
|
soup = BeautifulSoup(body, 'html.parser')
|
|
texts = soup.findAll(text=True)
|
|
visible_texts = filter(tag_visible, texts)
|
|
return u" ".join(t.strip() for t in visible_texts).splitlines()
|
|
|
|
###################################################################################################
|
|
# test a connection to an HTTP/HTTPS server
|
|
def test_connection(protocol="http", host="127.0.0.1", port=80, uri="", username=None, password=None, ssl_verify="full", user_agent="hedgehog"):
|
|
status = 400
|
|
message = "Connection error"
|
|
output = []
|
|
|
|
if protocol.lower() == "https":
|
|
if ssl_verify.lower() == "full":
|
|
c = HTTPSConnection(host, port=port)
|
|
else:
|
|
c = HTTPSConnection(host, port=port, context=ssl._create_unverified_context())
|
|
elif protocol.lower() == "http":
|
|
c = HTTPConnection(host)
|
|
else:
|
|
c = None
|
|
|
|
if c:
|
|
try:
|
|
if username and password:
|
|
c.request('GET', f'/{uri}', headers={ 'User-agent': user_agent, 'Authorization' : 'Basic %s' % b64encode(f"{username}:{password}".encode()).decode("ascii") })
|
|
else:
|
|
c.request('GET', f'/{uri}', headers={ 'User-agent': user_agent })
|
|
res = c.getresponse()
|
|
status = res.status
|
|
message = res.reason
|
|
output = text_from_html(res.read())
|
|
|
|
except Exception as e:
|
|
if len(output) == 0:
|
|
output = ["Error: {}".format(e)]
|
|
|
|
return status, message, output
|
|
|
|
###################################################################################################
|
|
# test if a remote port is open
|
|
def check_socket(host, port):
|
|
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
|
|
sock.settimeout(10)
|
|
if sock.connect_ex((host, port)) == 0:
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
###################################################################################################
|
|
# determine a list of available (non-virtual) adapters (Iface's)
|
|
def get_available_adapters():
|
|
|
|
available_adapters = []
|
|
_, all_iface_list = run_process("find /sys/class/net/ -mindepth 1 -maxdepth 1 -type l -printf '%P %l\\n'")
|
|
available_iface_list = [x.split(" ", 1)[0] for x in all_iface_list if 'virtual' not in x]
|
|
|
|
# for each adapter, determine its MAC address and link speed
|
|
for adapter in available_iface_list:
|
|
mac_address = '??:??:??:??:??:??'
|
|
speed = '?'
|
|
try:
|
|
with open(f"/sys/class/net/{adapter}/address", 'r') as f:
|
|
mac_address = f.readline().strip()
|
|
except:
|
|
pass
|
|
try:
|
|
with open(f"/sys/class/net/{adapter}/speed", 'r') as f:
|
|
speed = f.readline().strip()
|
|
except:
|
|
pass
|
|
description = f"{mac_address} ({speed} Mbits/sec)"
|
|
iface = Iface(adapter, description)
|
|
available_adapters.append(iface)
|
|
|
|
return available_adapters
|
|
|
|
###################################################################################################
|
|
# identify the specified adapter using ethtool --identify
|
|
def identify_adapter(adapter, duration=NIC_BLINK_SECONDS, background=False):
|
|
if background:
|
|
subprocess.Popen(["/sbin/ethtool", "--identify", adapter, str(duration)], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
else:
|
|
retCode, _ = run_process(f"/sbin/ethtool --identify {adapter} {duration}", stdout=False, stderr=False, timeout=duration*2)
|
|
return (retCode == 0)
|
|
|
|
###################################################################################################
|
|
# client that writes to the local instance of protologbeat listening on the configured host/port/protocol
|
|
class HeatBeatLogger:
|
|
|
|
def __init__(self, host='127.0.0.1', port=9515, proto='udp', format='plain', debug=False):
|
|
self.host = host
|
|
self.port = port
|
|
if proto == 'udp':
|
|
self.proto = 'udp'
|
|
self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
else:
|
|
self.proto = 'tcp'
|
|
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
self.format = format
|
|
if self.format not in ['plain','json']:
|
|
self.format = 'plain'
|
|
self.debug = debug
|
|
if self.debug:
|
|
print("Creating instance of logger via {} on {}:{}".format(self.proto, self.host, self.port))
|
|
|
|
def enable_debug(self):
|
|
self.debug = True
|
|
|
|
def send_message(self, msg):
|
|
if self.format == 'json':
|
|
payload = json.dumps(msg)
|
|
else:
|
|
payload = msg
|
|
if self.debug:
|
|
print("Sending message: {}".format(payload.encode('utf-8')))
|
|
self.socket.sendto(payload.encode('utf-8'), (self.host, self.port))
|