statscollect/collect.py

import subprocess
import time
import re
import csv
import os
import shutil
import socket
import platform
import tarfile

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 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", "-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("[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]
        print(f"mem: {lineData}")

        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", "15", "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):
    if is_freebsd():
        gstat_command = ["gstat", "-C", "-s", "-d", "-o", "-p", "-I", "5s"]
        with open("gstat.csv", "a") as output_file:
            # Create a Popen object with stdout redirected to the file
            process = subprocess.Popen(gstat_command, stdout=output_file)

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

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

        zpoolIostat()
        ifstat()

        if is_debian():
            scaleIostat_cpu()
            scaleIostat_disk()
            scaleMemstat()

        if is_freebsd():
            coreCPUstat()
            coreMemstat()

        if i == minutes:
            print("")
            break

        time.sleep(60)

    # kill gstat if freebsd
    if is_freebsd():
        process.kill()


def run_debug():
    print("Taking new 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 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 runs for x minutes and collects cpu/disk/network stats   #")
    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()

    minutes = int(input("Enter the duration in minutes: "))

    print("Starting Collection")

    # Collect data
    collect_data(minutes)

    # 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()

    print("Data collection and upload completed.")


if __name__ == "__main__":
    main()