#!/usr/bin/env python
# coding: utf-8

# # Title: msticpy - nbwidgets
# ## Description:
# This contains a few aggregated widgets using IPyWidgets that help speed things up during an investigation.
# 
# You must have msticpy installed to run this notebook:
# ```
# %pip install --upgrade msticpy
# ```
# 
# MSTICpy versions >= 0.8.5

# <a id='contents'></a>
# ## Table of Contents
# - [Setting query start/end times](#QueryTime)
# - [Simple time range](#Lookback)
# - [Selecting and Displaying Alerts](#AlertSelector)
# - [Selecting from list or dict](#SelectString)
# - [Getting a value from environment](#GetEnvironmentKey)
# 

# In[1]:


# Imports
import sys
MIN_REQ_PYTHON = (3,6)
if sys.version_info < MIN_REQ_PYTHON:
    print('Check the Kernel->Change Kernel menu and ensure that Python 3.6')
    print('or later is selected as the active kernel.')
    sys.exit("Python %s.%s or later is required.\n" % MIN_REQ_PYTHON)

from IPython.display import display, Markdown
import pandas as pd
# Import nbtools package
from msticpy.nbtools import *


# <a id='QueryTime'></a>[Contents](#contents)
# ## QueryTime
# 
# This widget is used to specify time boundaries - designed to be used with the built-in msticpy queries and custom queries.
# The `start` and `end` times are exposed as datetime properties.
# 
# ```
# QueryTime.
# 
# Composite widget to capture date and time origin
# and set start and end times for queries.
# 
# Parameters
# ----------
# QueryParamProvider : QueryParamProvider
#     Abstract base class
# 
# Parameters
# ----------
# origin_time : datetime, optional
#     The origin time (the default is `datetime.utcnow()`)
# label : str, optional
#     The description to display
#     (the default is 'Select time ({units}) to look back')
# before : int, optional
#     The default number of `units` before the `origin_time`
#     (the default is 60)
# after : int, optional
#     The default number of `units` after the `origin_time`
#     (the default is 10)
# max_before : int, optional
#     The largest value for `before` (the default is 600)
# max_after : int, optional
#     The largest value for `after` (the default is 100)
# units : str, optional
#     Time unit (the default is 'min')
#     Permissable values are 'day', 'hour', 'minute', 'second'
#     These can all be abbreviated down to initial characters
#     ('d', 'm', etc.)
# auto_display : bool, optional
#     Whether to display on instantiation (the default is False)
# ```

# In[2]:


q_times = nbwidgets.QueryTime(units='day', max_before=20, before=5, max_after=1)
q_times.display()


# In[3]:


print(q_times.start, '....', q_times.end)


# Keep multiple query boundaries aligged by having QueryTime instances reference the time of the same alert or event, or have them chained from one another by referencing the origin_time of an earlier QueryTimes object

# In[4]:


from datetime import datetime, timedelta
class MyAlert:
    pass
alert = MyAlert()
alert.TimeGenerated = datetime.utcnow() - timedelta(15)
alert.TimeGenerated

q_times1 = nbwidgets.QueryTime(units='hour', max_before=20, before=1, max_after=1, 
                             origin_time=alert.TimeGenerated, auto_display=True)

q_times2 = nbwidgets.QueryTime(units='hour', max_before=20, before=4, max_after=2, 
                             origin_time=alert.TimeGenerated, auto_display=True)


# In[5]:


alert.TimeGenerated = datetime.utcnow()
q_times1 = nbwidgets.QueryTime(units='hour', max_before=20, before=1, max_after=1, 
                             origin_time=alert.TimeGenerated, auto_display=True)

q_times2 = nbwidgets.QueryTime(units='hour', max_before=20, before=4, max_after=2, 
                             origin_time=q_times2.origin_time, auto_display=True)


# In[6]:


# Use the values in a query
my_kql = f'''
SecurityAlert 
| where TimeGenerated >= datetime({q_times1.start})
| where TimeGenerated <= datetime({q_times1.end})'''
print(my_kql)


# <a id='Lookback'></a>[Contents](#contents)
# ## Lookback
# Simpler version with single slider value
# 
# Docstring:
# `nbtools.Lookback?`

# In[7]:


alert.TimeGenerated = datetime.utcnow() - timedelta(5)
lb = nbwidgets.Lookback(origin_time=alert.TimeGenerated, auto_display=True, max_value=48)


# In[8]:


print(lb.start, '....', lb.end)


# <a id='AlertSelector'></a>[Contents](#contents)
# ## Alert Browser
# 
# ```
# SelectAlert.
# 
# View list of alerts and select one for investigation.
# Optionally provide and action to call with the selected alert as a parameter
# (typically used to display the alert.)
# 
# Attributes:
#     selected_alert: the selected alert
#     alert_id: the ID of the selected alert
#     alerts: the current alert list (DataFrame)
# Init docstring:
# Create a new instance of AlertSelector.
# 
# Parameters
# ----------
# alerts : pd.DataFrame
#     DataFrame of alerts.
# action : Callable[..., None], optional
#     Optional function to execute for each selected alert.
#     (the default is None)
# columns : list, optional
#     Override the default column names to use from `alerts`
#     (the default is ['StartTimeUtc', 'AlertName',
#     'CompromisedEntity', 'SystemAlertId'])
# auto_display : bool, optional
#     Whether to display on instantiation (the default is False)
# ```

# ### Simple alert selector
# Selected alert is available as `select_alert_widget.selected_alert` property

# In[9]:


# Load test data
alerts = pd.read_csv('data/alertlist.csv')

display(Markdown('### Simple alert selector'))
display(Markdown('Selected alert is available as `select_alert_widget.selected_alert`'))
alert_select = nbwidgets.SelectAlert(alerts=alerts)
alert_select.display()


# ### Alert selector with action=SecurityAlert'
# You can pass a function that returns one or more displayable objects.
# You can also pass a class (in this case we're passing `SecurityAlert`) that produces an IPython displayable object.
# 
# The `action` class/function is passed the raw alert row as a parameter, as it is selected from the list

# In[10]:


alert_select = nbwidgets.SelectAlert(alerts=alerts, action=SecurityAlert)
alert_select.display()


# In[11]:


from IPython.display import HTML
security_alert = None

# create a function to get the displayable object
def alert_with_entities(alert):
    return HTML(SecurityAlert(alert).to_html(show_entities=True))
                
alert_select = nbwidgets.SelectAlert(alerts=alerts.query('CompromisedEntity == "MSTICALERTSWIN1"'), 
                                     action=alert_with_entities)
display(Markdown('### Or a more detailed display with extracted entities'))
alert_select


# <a id='SelectItem'></a>[Contents](#contents)
# ## SelectItem
# 
# Similar to AlertSelector but simpler and allows you to use any list or dictionary of items.
# 
# ```
# Selection list from list or dict.
# 
# Attributes:
#     value : The selected value.
# Init docstring:
# Select an item from a list or dict.
# 
# Parameters
# ----------
# description : str, optional
#     The widget label to display (the default is None)
# item_list : List[str], optional
#     A `list` of items to select from (the default is None)
# item_dict : Mapping[str, str], optional
#     A `dict` of items to select from. When using `item_dict`
#     the keys are displayed as the selectable items and value
#     corresponding to the selected key is set as the `value`
#     property.
#     (the default is None)
# action : Callable[..., None], optional
#     function to call when item selected (passed a single
#     parameter - the value of the currently selected item)
#     (the default is None)
# auto_display : bool, optional
#     Whether to display on instantiation (the default is False)
# height : str, optional
#     Selection list height (the default is '100px')
# width : str, optional
#     Selection list width (the default is '50%')
# ```

# In[12]:


# extract the entities from the previously selected alert
security_alert = SecurityAlert(alert_select.selected_alert)
if security_alert is None:
    security_alert = SecurityAlert(alerts.iloc[1])
ent_dict = {ent['Type']:ent for ent in security_alert.entities}

# from IPython.display import HTML

# # create a display function for the entities
# def entity_to_html(entity):
#     e_text = str(entity)
#     e_type = entity.Type
#     e_text = e_text.replace("\n", "<br>").replace(" ", "&nbsp;")
#     return HTML(f"<h3>{e_type}</h3>{e_text}")
    
nbwidgets.SelectItem(item_dict=ent_dict,
                       description='Select an item',
                       action=lambda x: x,
                       auto_display=True);


# <a id='GetEnvironmentKey'></a>[Contents](#contents)
# ## GetEnvironmentKey
# Get editable value of environment variable. Common use would be retrieving an API key from your environment or allowing you to paste in a value if the environment key isn't set.
# 
# Note setting the variable only persists in the python kernel process running at the time. So you can retrieve it later in the notebook but not in other processes.

# In[13]:


nbwidgets.GetEnvironmentKey(env_var='userprofile', auto_display=True);


# <a id='SelectSubset'></a>[Contents](#contents)
# ## SelectSubset
# Allows you to select one or multiple items from a list to populate an output set.
# 
# ```
# Class to select a subset from an input list.
# 
#     Attributes
#     ----------
#     selected_values : List[Any]
#         The selected item values.
#     selected_items : List[Any]
#         The selected items label and value
#         
# Init docstring:
# Create instance of SelectSubset widget.
# 
# Parameters
# ----------
# source_items : Union[Dict[str, str], List[Any]]
#     List of source items - either a dictionary(label, value),
#     a simple list or
#     a list of (label, value) tuples.
# default_selected : Union[Dict[str, str], List[Any]]
#     Populate the selected list with values - either
#     a dictionary(label, value),
#     a simple list or
#     a list of (label, value) tuples.
# ```

# In[14]:


# Simple list
items = list(alerts["AlertName"].values)
sel_sub = nbwidgets.SelectSubset(source_items=items)


# In[15]:


# Label/Value pair items with a a subset of pre-selected items
items = {v: k for k, v in alerts["AlertName"].to_dict().items()}
pre_selected = {v: k for k, v in alerts["AlertName"].to_dict().items() if "commandline" in v}
sel_sub = nbwidgets.SelectSubset(source_items=items, default_selected=pre_selected)


# In[16]:


print("Values:", sel_sub.selected_values, "\n")
print("Items:", sel_sub.selected_items)


# ## Progress Indicator

# In[17]:


from time import sleep
progress = nbwidgets.Progress(completed_len=2000)
for i in range(0, 2100, 100):
    progress.update_progress(new_total=i)
    sleep(0.1)
    
inc_progress = nbwidgets.Progress(completed_len=10)
for i in range(0, 11):
    inc_progress.update_progress(delta=1)
    sleep(0.1)
print("Volume goes to eleven!")


# ## Logon Display
# Display logon details for a Windows or Linux logon

# In[18]:


win_logons = pd.read_csv("data/host_logons.csv")
user_dict = win_logons.apply(lambda x: f"{x.TargetDomainName}/{x.TargetUserName}   ({x.TimeGenerated})", axis=1).to_dict()
user_dict = {v: k for k, v in user_dict.items()}

from msticpy.nbtools.nbdisplay import format_logon
# create a display function for the entities
def disp_logon(index):
    print
    logons = win_logons[win_logons.index == index]
    return format_logon(logons)
    
acct_select = nbwidgets.SelectItem(item_dict=user_dict,
                       description='Select an item',
                       action=disp_logon,
                       auto_display=True);


# #### Display a list of logons

# In[19]:


# display a list of logons
display(format_logon(win_logons.head(5)))


# ## Registered Widgets
# 
# Some of the widgets (QueryTimes, GetText) can register themselves and retain
# the setting and values previously entered. This can be useful when stepping through
# a notebook since it is a common mistake to enter text in a text box and then
# execute the same cell again by mistake. This, of course, usually results in the 
# widget being reset to its default state and erasing the values you just entered.
# 
# If you use a registered widget and then create a new copy of the widget with identical
# parameters it will look in the registry for a previous copy of itself and auto-populate
# it's values with the previous-entered ones.
# 
# Registered widgets can also read their default values from notebook variables - this
# is mainly useful with notebooks that are programmatically supplied with 
# parameters and executed with something like Papermill.
# 
# Several of the additional parameters available in RegisteredWidgets init are
# for internal use by widgets but three are usable by users:
# ```
#         Parameters
#         ----------
#         nb_params : Optional[Dict[str, str]], optional
#             A dictionary of attribute names and global variables. If the variable
#             exists in the global namespace it will be used to populate the
#             corresponding widget attribute. This is only done if the widget
#             attribute currently has no value (i.e. restoring a value from
#             the registry takes priority over this),
#             by default None
#         ns : Dict[str, Any], optional
#             Namespace to look for global variables, by default None
#         register : bool
#             Do not register the widget or retrieve values from previously-
#             registered instance.
# ```

# In[20]:


print(nbwidgets.RegisteredWidget.__init__.__doc__)


# In[21]:


mem_text = nbwidgets.GetText(prompt="Enter your name")

# we insert a value here to mimic typing something in the text box
mem_text._value = "Ian"
mem_text


# When we re-execute the cell or use the same widget with identical arguments
# the value is populated from the registry cache

# In[22]:


mem_text = nbwidgets.GetText(prompt="Enter your name")
mem_text


# #### QueryTime also supports registration

# In[23]:


from datetime import datetime, timedelta
q_times = nbwidgets.QueryTime(auto_display=True, max_before=12, max_after=2, units="day")


# In[24]:


# mimic setting values in the control (these don't update the display)
q_times.origin_time = datetime.utcnow() - timedelta(5)
q_times.before = 3
q_times.after = 5


# Note the origin, before and after have all been copied from the previous instance

# In[25]:


q_times = nbwidgets.QueryTime(auto_display=True, max_before=12, max_after=2, units="day")


# ### To skip registration add the parameter `register=False`

# In[26]:


q_times = nbwidgets.QueryTime(auto_display=True, max_before=12, max_after=2, units="day", register=False)


# ### Using notebook parameters to populate RegisteredWidgets

# In[27]:


# This might be defined in a parameter cell at the beginning of the noteboook
my_name = "The other Ian"

my_text = nbwidgets.GetText(prompt="enter your real name", nb_params={"_value": "my_name"}, ns=globals())
my_text


# ## Multi-Option buttons with async wait
# This widget is pretty simple on the surface but has some useful features
# for waiting for user input.
# 

# In[28]:


opt = nbwidgets.OptionButtons(
    description="Do you really want to do this?",
    buttons=["Confirm", "Skip", "Cancel"]
)

# Displaying the widget works as expected
# and sets `widget.value` to the last chosen button value.
opt


# ### Using OptionButtons to wait until an option is chosen (or timeout expires)
# Option buttons uses an asynchronous event loop to track both the button
# state and the timeout simultaneously.
# 
# Because this requires the use of asynchronous code you must do the following
# - call *widget*`.display_async()` method rather than just `display()` or using the auto-display functionality of Jupyter
# - prefix this call with `await` - this tells IPython/Jupyter that you are executing asynchronous code and that it needs
#   to wait until this call has completed before continuing with cell execution. 

# In[29]:


# Using display_async will run the widget with a visible
# timer. As soon as one option is chosen, that remains as the value
# of the value of the widget.value property.
opt = nbwidgets.OptionButtons(description="Continue?", timeout=10)
await opt.display_async()


# > **Note**
# > Awaiting the OptionButtons control does not pause the notebook execution.
# > This is a capability that we are still working on.