def add_pivot_funcs(pivot: Pivot = None, **kwargs):
"""
Add notebooklet run functions as pivot methods.
Parameters
----------
pivot : Pivot, optional
Pivot instance.
kwargs : Dict[str, Any]
Optional keyword arguments to pass to Pivot init.
"""
if not pivot:
piv_kwargs = {
key: arg
for key, arg in kwargs.items() if key in ("namespace", "providers")
}
pivot = Pivot.current or Pivot(**piv_kwargs)
for nb_name, nb_class in nblts.iter_classes():
if not issubclass(nb_class, Notebooklet) or nb_name == "TemplateNB":
continue
if "value" not in nb_class.metadata.inputs:
# This doesn't take a "value" input so can't use as a pivot
continue
nb_obj = nb_class()
run_func = getattr(nb_obj, "run")
wrp_func = _wrap_run_func(run_func, pivot.get_timespan)
func_new_name = _to_py_name(nb_name)
entity_map: Dict[str, str] = {}
for entity in nb_class.metadata.entity_types:
if entity not in _ENTITY_MAP:
continue
entity_map.update(_ENTITY_MAP[entity])
if not entity_map:
continue
piv_reg = PivotRegistration(
input_type="value",
entity_map=entity_map,
func_new_name=func_new_name,
src_func_name="run",
can_iterate=False,
func_input_value_arg="value",
return_raw_output=True,
)
Pivot.add_pivot_function(func=wrp_func,
pivot_reg=piv_reg,
container="nblt")
def _init_pivot(monkeypatch):
test_data = str(Path(TEST_DATA_PATH).absolute())
monkeypatch.setattr(data_providers, "GeoLiteLookup", GeoIPLiteMock)
data_providers.init(
query_provider="LocalData",
providers=["geolitelookup"],
LocalData_data_paths=[test_data],
LocalData_query_paths=[test_data],
)
return Pivot()
def test_pivot_time(data_providers):
"""Function_docstring."""
providers = data_providers.values()
end = datetime.utcnow()
start = end - timedelta(1)
timespan = TimeSpan(start=start, end=end)
pivot = Pivot(providers=providers, timespan=timespan)
check.equal(pivot.start, start)
check.equal(pivot.end, end)
end = end - timedelta(1)
start = start - timedelta(1)
timespan = TimeSpan(start=start, end=end)
pivot.timespan = timespan
check.equal(pivot.start, start)
check.equal(pivot.end, end)
_fake_provider_connected(data_providers["az_sent_prov"])
query = entities.Host.AzureSentinel.list_host_processes(host_name="test",
print_query=True)
check.is_in(start.isoformat(), query)
check.is_in(end.isoformat(), query)
def test_pivot_time(data_providers):
"""Function_docstring."""
providers = data_providers.values()
end = datetime.utcnow()
start = end - timedelta(1)
timespan = TimeSpan(start=start, end=end)
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=UserWarning)
pivot = Pivot(providers=providers, timespan=timespan)
check.equal(pivot.start, start)
check.equal(pivot.end, end)
end = end - timedelta(1)
start = start - timedelta(1)
# Test different ways of setting the time
timespan = TimeSpan(start=start, end=end)
pivot.timespan = timespan
check.equal(pivot.start, start)
check.equal(pivot.end, end)
pivot.timespan = _TimeObj(start=timespan.start, end=timespan.end)
check.equal(pivot.start, start)
check.equal(pivot.end, end)
pivot.set_timespan(timespan)
check.equal(pivot.start, start)
check.equal(pivot.end, end)
pivot.set_timespan(start=timespan.start, end=timespan.end)
check.equal(pivot.start, start)
check.equal(pivot.end, end)
# Make sure the values provided to queries match.
_fake_provider_connected(data_providers["az_sent_prov"])
query = entities.Host.AzureSentinel.SecurityEvent_list_host_processes(
host_name="test", print=True)
check.is_in(start.isoformat(), query)
check.is_in(end.isoformat(), query)
def test_pivot_shortcuts():
"""Test pivot function shortcut creation and deletion."""
Pivot()
check.is_true(hasattr(IpAddress, "util"))
util_ctnr = getattr(IpAddress, "util")
func = getattr(util_ctnr, "ip_type")
IpAddress.make_pivot_shortcut("util.ip_type", "test_iptype")
check.is_true(hasattr(IpAddress, "test_iptype"))
check.equal(func, IpAddress.test_iptype)
ip_addr = IpAddress(Address="192.168.1.2")
ip_df = ip_addr.test_iptype()
check.is_instance(ip_df, pd.DataFrame)
with pytest.raises(AttributeError):
IpAddress.make_pivot_shortcut("util.not_defined", "test_iptype")
with pytest.raises(TypeError):
IpAddress.make_pivot_shortcut("properties", "test_iptype")
with pytest.raises(AttributeError):
IpAddress.make_pivot_shortcut("util.ip_type", "test_iptype")
IpAddress.make_pivot_shortcut("util.ip_type",
"test_iptype",
overwrite=True)
check.is_true(hasattr(IpAddress, "test_iptype"))
check.equal(func, IpAddress.test_iptype)
IpAddress.del_pivot_shortcut("test_iptype")
check.is_false(hasattr(IpAddress, "test_iptype"))
with pytest.raises(AttributeError):
IpAddress.del_pivot_shortcut("test_iptype")
with pytest.raises(TypeError):
IpAddress.del_pivot_shortcut("properties")
def test_remove_pivots(_create_pivot_ns):
"""Test remove pivots function."""
piv_attrs = _get_piv_attrs(entities.Host)
check.is_true(piv_attrs)
with pytest.raises(ValueError):
Pivot.remove_pivot_funcs(entity="TestEntity")
piv_attrs = _get_piv_attrs(entities.Host)
check.is_true(piv_attrs)
Pivot.remove_pivot_funcs(entity="Host")
piv_attrs = _get_piv_attrs(entities.Host)
check.is_false(piv_attrs)
piv_attrs = _get_piv_attrs(entities.IpAddress)
check.is_true(piv_attrs)
Pivot.remove_pivot_funcs(entity="all")
piv_attrs = _get_piv_attrs(entities.IpAddress)
check.is_false(piv_attrs)
def _create_pivot(data_providers):
_reset_entities()
providers = data_providers.values()
return Pivot(providers=providers)
def _create_pivot_ns(data_providers):
_reset_entities()
locals().update(data_providers)
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=UserWarning)
return Pivot(namespace=locals())
def _create_pivot_list(data_providers):
_reset_entities()
providers = data_providers.values()
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=UserWarning)
return Pivot(providers=providers)
def _reset_entities():
"""Clear any query containers in entities."""
Pivot.remove_pivot_funcs(entity="all")
def _create_pivot():
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=UserWarning)
return Pivot()
def _create_pivot_ns(data_providers):
_reset_entities()
locals().update(data_providers)
return Pivot(namespace=locals())