statscollect/collect.py

import subprocess
import time
import re
import csv
import os
import shutil
import socket
import platform
import argparse
import logging
import glob

from datetime import datetime


def is_freebsd():
    return platform.system() == "FreeBSD"


def is_debian():
    # Debian-based systems often have 'debian' in their platform string
    return "linux" in platform.platform().lower()


def getTimestamp():
    timestamp = time.time()
    local_time = time.localtime(timestamp)
    # Format the time components for a more readable output
    return time.strftime("%Y-%m-%d %H:%M:%S", local_time)


def runCollect(command):
    # Run the command and capture output
    result = subprocess.run(command, capture_output=True, text=True)
    # Check the return code (0 for success)
    if result.returncode == 0:
        # Access the captured output as a string
        return result.stdout

    return result.stderr



def isPoolScrubbing():

    command="zpool status"

    grep_string = "scrub.in.progress"
    
    try:
        process = subprocess.Popen(
            command,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,  # Important for returning strings rather than bytes
        )
        stdout, stderr = process.communicate()

        if process.returncode == 0:  # Command executed successfully
            if re.search(grep_string, stdout, re.MULTILINE).group():
                print("WARNING: Pool is SCRUBBING")
                return True
            else:
                print("Checked for Pool Scrubbing;  Good to Go...")
                return False
        else:
            # Command failed. Print stderr for debugging.
            print(f"Error executing command: {command}")
            print(stderr) #Print the error that the command returned.
            exit()

    except FileNotFoundError:
        print(f"Command not found: {command}")
        exit()
    except Exception as e:
        print(f"An error occurred: {e}")
        exit()


def coreIostat_disk():
    command = ["iostat", "-xd", "-t", "da", "1", "1"]
    filename = "ioStat.csv"
    collect = runCollect(command)

    if collect:
        byline = re.split("\n", collect.strip())
        data = byline[2:]

        for line in data:
            lineData = line.split()

            if lineData:
                with open(filename, "a", newline="") as csvfile:
                    csv_writer = csv.writer(csvfile)
                    lineData.insert(0, getTimestamp())
                    csv_writer.writerow(lineData)

    else:
        print(f"Error running command: {collect}")


def scaleIostat_disk():
    command = ["iostat", "-xyd", "1", "1"]
    filename = "ioStat.csv"
    collect = runCollect(command)

    if collect:
        byline = re.split("\n", collect.strip())
        data = byline[4:]

        for line in data:
            lineData = line.split()

            if lineData:
                with open(filename, "a", newline="") as csvfile:
                    csv_writer = csv.writer(csvfile)
                    lineData.insert(0, getTimestamp())
                    csv_writer.writerow(lineData)

    else:
        print(f"Error running command: {collect}")


def scaleIostat_cpu():
    command = ["iostat", "-c", "1", "1"]
    filename = "cpuStat.csv"

    collect = runCollect(command)

    if collect:
        byline = re.split("\n", collect.strip())
        data = byline[3].strip()
        lineData = data.split()

        if lineData:
            with open(filename, "a", newline="") as csvfile:
                csv_writer = csv.writer(csvfile)
                lineData.insert(0, "cpu")
                lineData.insert(0, getTimestamp())
                csv_writer.writerow(lineData)

    else:
        print(f"Error running command: {collect}")


def convert_to_megabytes(data_size):
    try:
        unit = data_size[-1].upper()
        size = data_size[:-1]
        size = float(size)
    except (ValueError, AttributeError):
        return None
    unit_map = {
        "K": 1 / (1024**1),
        "M": 1,
        "G": 1024,
        "T": 1024**2,
    }
    if unit not in unit_map:
        return None
    return int(size * unit_map[unit])


def scaleMemstat():
    command = ["top", "-bn", "1"]

    filename = "memStat.csv"
    collect = runCollect(command)

    if collect:
        mem = re.search("Mem :.*", collect).group()
        mem = re.sub("Mem..", "", mem.strip())
        mem = re.sub("[total|free|used|buff|cache|,|/]", "", mem.strip()).split()

        meminmeg = []

        for x in mem:
            meminmeg.append(convert_to_megabytes(f"{x}M"))

        mem = meminmeg

        mem_free = mem[1]
        mem_used = mem[2]

        lineData = [mem_used, mem_free]

        if lineData:
            with open(filename, "a", newline="") as csvfile:
                csv_writer = csv.writer(csvfile)
                lineData.insert(0, "mem")
                lineData.insert(0, getTimestamp())
                csv_writer.writerow(lineData)
    else:
        print(f"Error running command: {collect}")


def coreMemstat():
    command = ["top", "-n", "1"]

    filename = "memStat.csv"
    collect = runCollect(command)

    if collect:
        mem = re.search("Mem:.*", collect).group()
        mem = re.sub("Mem.", "", mem.strip())
        mem = re.sub("[Active|Inact|Wired|Free|,]", "", mem.strip()).split()

        meminmeg = []

        for x in mem:
            meminmeg.append(convert_to_megabytes(x))

        mem = meminmeg

        mem_used = mem[0]
        mem_free = mem[3]

        lineData = [mem_used, mem_free]

        if lineData:
            with open(filename, "a", newline="") as csvfile:
                csv_writer = csv.writer(csvfile)
                lineData.insert(0, "mem")
                lineData.insert(0, getTimestamp())
                csv_writer.writerow(lineData)
    else:
        print(f"Error running command: {collect}")


def ifstat():
    command = ["ifstat", "-znq", "1", "1"]

    filename = "ifStat.csv"
    collect = runCollect(command)
    if collect:
        # process collection string
        bylines = collect.split("\n")
        interfaces = bylines[0].split()
        stats = bylines[2].split()

        for nic in interfaces:
            lineData = [getTimestamp(), nic, stats.pop(0), stats.pop(0)]
            if lineData:
                with open(filename, "a", newline="") as csvfile:
                    csv_writer = csv.writer(csvfile)
                    csv_writer.writerow(lineData)
    else:
        print(f"Error running command: {collect}")


def coreCPUstat():
    command = ["iostat", "-C", "-t", "proc", "-d"]
    filename = "cpuStat.csv"

    collect = runCollect(command)

    if collect:
        # process collection string
        bylines = collect.split("\n")
        lineData = bylines[2].split()

        lineData = lineData[:-2] + [0] + lineData[-2:]

        if lineData:
            with open(filename, "a", newline="") as csvfile:
                csv_writer = csv.writer(csvfile)
                lineData.insert(0, "cpu")
                lineData.insert(0, getTimestamp())
                csv_writer.writerow(lineData)

    else:
        print(f"Error running command: {collect}")


def zpoolIostat():
    command = ["zpool", "iostat", "-Tu", "-l", "-p", "-v", "-y", "5", "1"]
    collect = runCollect(command)

    if collect:
        # split output by pools
        pools = re.split(r"^-", collect, flags=re.MULTILINE)

        i = 0
        for pool in pools:
            if i == 0:
                pass
            else:
                # remove ------ lines
                pool = re.sub("^---.*", "", pool).strip()

                # turn - values into 0
                pool = re.sub(" -", "0", pool)

                y = 0

                poolbyline = re.split("\n", pool)

                for line in poolbyline:
                    lineData = line.split()

                    if not lineData:
                        break

                    if y == 0:
                        poolname = lineData[0]
                        filename = poolname + "-zio.csv"

                    # Now open the file in append mode ('a')
                    with open(filename, "a", newline="") as csvfile:
                        csv_writer = csv.writer(csvfile)
                        lineData.insert(0, getTimestamp())
                        csv_writer.writerow(lineData)

                    y += 1

            i += 1

    else:
        print(f"Error running command: {collect}")


def collect_data(minutes):
    for i in range(minutes):
        if i == 0:
            print("Minute:", end="", flush=True)

        print(f" {i}", end="", flush=True)

        logging.info("running zpool iostat")
        zpoolIostat()
        logging.info("running ifstat")
        ifstat()

        if is_debian():
            logging.info("running iostat_cpu")
            scaleIostat_cpu()
            logging.info("running iostat_disk")
            scaleIostat_disk()
            logging.info("running memstat")
            scaleMemstat()

        if is_freebsd():
            logging.info("running cpustat")
            coreCPUstat()
            logging.info("running memstat")
            coreMemstat()

            # gives same data each interval
            # coreIostat_disk()

        if i == minutes:
            print("")
            print("")
            logging.info("collection over")
            break

        time.sleep(55)


def run_debug():
    print("Taking new debug.")
    logging.info("taking debug")
    command = ["midclt", "call", "system.debug_generate", "-job"]
    try:
        result = subprocess.run(command, capture_output=True, text=True, check=True)
        return result.stdout
    except subprocess.CalledProcessError as e:
        print(f"Error running midclt debug: {e}")
        return None


def cleanCSV():
    if is_debian():
        dir = "/var/log/proftpd"

    if is_freebsd():
        dir = "/var/log"

    csv_files = glob.glob(os.path.join(dir, "*.csv"))

    for filename in csv_files:
        os.remove(filename)


def collect_csv():
    source_dir = os.getcwd()

    if is_debian():
        destination_dir = "/var/log/proftpd"

    if is_freebsd():
        destination_dir = "/var/log"

    for filename in os.listdir(source_dir):
        if filename.endswith(".csv"):  # Check for .csv extension
            source_file = os.path.join(source_dir, filename)
            destination_file = os.path.join(destination_dir, filename)

            # Handle potential errors (e.g., source file not found, permission issues)
            try:
                shutil.copy2(source_file, destination_file)
                print(f"Copied '{filename}' to /var/log successfully.")
            except FileNotFoundError:
                print(f"Error: File '{filename}' not found in '{source_dir}'.")
            except PermissionError:
                print(f"Error: Insufficient permissions to copy '{filename}'.")
            except Exception as e:  # Catch other potential errors
                print(f"Error copying '{filename}': {e}")


def upload_debug():
    """
    Uploads debug file to a specified FTP server.
    """

    hostname = socket.gethostname()

    # Generate debug file
    debug_file = run_debug().strip()

    # Upload data files
    subprocess.run(
        [
            "curl",
            "--user",
            "customer:ixcustomer",
            "-T",
            debug_file,
            f"ftp.ixsystems.com/debug-perf-{hostname}-{datetime.now().strftime('%Y%m%d%H%M')}.tgz",
        ]
    )


def welcome():
    print("########################################################################")
    print("#         FreeNAS CORE/SCALE performance capture script v.02           #")
    print("########################################################################")
    print("#         This script collects cpu/mem/disk/network/pool               #")
    print("# When it is completed, it will copy csv files to your /var/log folder #")
    print("# It will then attempt to take a debug and upload both to our FTP      #")
    print("# If not connected to the internet, it will have to be manually d/l    #")
    print("########################################################################")
    print("#   Running this script repeatedly will append results to CSV files    #")
    print("#       - https://gitlab.komputernerds.com/mmance/statscollect -       #")
    print("########################################################################")
    print("")


def main():
    welcome()
    isPoolScrubbing()
    
    logging.basicConfig(
        filename="perf.log",
        level=logging.DEBUG,
        format="%(asctime)s %(levelname)s %(message)s",
    )

    # setup argparse
    parser = argparse.ArgumentParser(description="TrueNAS Performance Capture Script")
    parser.add_argument(
        "runMinutes", nargs="?", help="Number of minutes to collect data.", default=0
    )
    parser.add_argument(
        "waitMinutes",
        nargs="?",
        help="Number of minutes to wait before collecting data.",
        default=0,
    )

    args = parser.parse_args()
    minutes = int(args.runMinutes)
    minutesToWait = int(args.waitMinutes)

    try:
        if not minutes:
            minutes = int(input("Enter the duration in minutes (1440 for 24hrs):"))
            minutesToWait = int(input("Enter the delay before capture in minutes: "))

        if not minutesToWait:
            minutesToWait = 0

        print(
            f"Set to wait {minutesToWait} minutes and then capture for {minutes} minutes."
        )

        logging.info(f"SETUP: {minutesToWait}m delay, {minutes}m of capture.")

        if minutesToWait:
            print(f"Delaying capture by {minutesToWait} minutes...")
            for i in range(minutesToWait):
                print(f"{i} ", end="", flush=True)
                time.sleep(minutesToWait * 60)
                print("")
    except KeyboardInterrupt:
        print("")
        logging.info("Caught control-c, exiting...")
        exit()

    try:
        print("Starting Collection")
        logging.info("Starting Collection")
        # Collect data
        if is_freebsd():
            # timeout = f"{minutes}m"
            #     "timeout",
            #     timeout,

            gstat_command = [
                "gstat",
                "-Csdop",
                "-I",
                "15s",
            ]

            logging.info("opening gstat.csv")
            gstatData = open("gstat.csv", "ab")
            logging.info("opening gstat.err")
            gstatError = open("gstat.err", "a")

            # with  as output_file:
            #     # Create a Popen object with stdout redirected to the file
            logging.info("starting Popen")
            process = subprocess.Popen(
                gstat_command, stdout=gstatData, stderr=gstatError
            )

        collect_data(minutes)

        # # kill gstat if freebsd
        if is_freebsd():
            logging.info("killing gstat")
            process.terminate()

            gstatData.flush()
            gstatError.flush()

            gstatData.close()
            gstatError.close()

    except KeyboardInterrupt:
        print("Caught Ctrl-C,  cancelling collection...")
        # kill gstat if freebsd
        if is_freebsd():
            logging.info("killing gstat")
            process.terminate()

            gstatData.flush()
            gstatError.flush()

            gstatData.close()
            gstatError.close()

        exit()

    # Copy data files to /var/log (replace with appropriate copying function)
    collect_csv()

    # Upload data (replace with credentials and actual upload logic if needed)
    upload_debug()
    cleanCSV()

    print("Data collection and upload completed.")


if __name__ == "__main__":
    main()