#!/usr/bin/env python3
# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""Simple utility to time the performance of the REST protocol."""

import sys
import timeit

from sigma.protocol import rest, xmlproto
from sigma.util import xmltime

TESTDATA = ("0009interface0002ip0003mac0009broadcast0007netmask0007gateway0010dhcpserver0011dnshostname\0\n"
            + 100*"0004eth0001110.69.33.14001200111146AA76001210.69.35.2550013255.255.252.0000000000031guido.esi.elementalsecurity.com\0\n")

NOW = xmltime.now()
TESTUPLOAD = xmlproto.Upload(type="interfaces.full",
                             startTime=NOW, stopTime=NOW)

commands = {
    "home": "obj, mtime = c.getAgentHome(mtime)",
    "home1": "obj, mtime = c.getAgentHome(None)",
    "dir": "obj, mtime = c.getDirective('5002.1', mtime)",
    "matrix": "obj, mtime = c.getMatrix(False, mtime)",
    "matrix1": "obj, mtime = c.getMatrix(False, None)",
    "upload": "assert c.uploadFile(TESTDATA, TESTUPLOAD)",
    "slow": "obj, mtime = rest.Connection(ssl=False, useTimeout=False).getAgentHome(mtime)",
    }

def timer(what, cmd):
    ssl = False
    count = 1000
    if what == "upload":
        ssl = True
        count = 50
    print("%-7s %5s %6d" % (what, ssl and "SSL" or "HTTP", count), end=" ")
    sys.stdout.flush()
    setup = ("from __main__ import rest, TESTDATA, TESTUPLOAD\n" +
             "c = rest.Connection(ssl=%r, useTimeout=False)\n" % ssl +
             "obj, mtime = c.getAgentHome(None)\n")
    t = timeit.Timer(cmd, setup)
    dt = t.timeit(count)
    print("%6.3f %9.3f  %s" % (dt, 1e3*dt/count, cmd))

if __name__ == "__main__":
    print("%-7s %5s %6s %6s %9s  %s" % ("Test", "Proto", "Calls",
                                        "Secs", "msec/Call", "Cmd"))
    keys = sorted(commands.keys())
    for key in keys:
        timer(key, commands[key])
