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

# # Title: msticpy - Folium Map Plotting
# 
# ## Introduction
# This module contains a class that wraps the `folium` package to plot geo-location data.
# 
# Read the [Folium documentation](https://python-visualization.github.io/folium/)
# 
# You must have msticpy installed to run this notebook:
# ```
# !pip install --upgrade msticpy
# ```

# <h1>Table of Contents<span class="tocSkip"></span></h1>
# <div class="toc"><ul class="toc-item"><li><span><a href="#Introduction" data-toc-modified-id="Introduction-1">Introduction</a></span></li><li><span><a href="#FoliumMap-class" data-toc-modified-id="FoliumMap-class-2">FoliumMap class</a></span></li><li><span><a href="#Adding-IP-Entities-to-the-map" data-toc-modified-id="Adding-IP-Entities-to-the-map-3">Adding IP Entities to the map</a></span><ul class="toc-item"><li><span><a href="#Plot-IPAddress-entities-with-location-data" data-toc-modified-id="Plot-IPAddress-entities-with-location-data-3.1">Plot IPAddress entities with location data</a></span></li><li><span><a href="#Use-different-colors-and-icons" data-toc-modified-id="Use-different-colors-and-icons-3.2">Use different colors and icons</a></span></li></ul></li><li><span><a href="#Custom-Icons" data-toc-modified-id="Custom-Icons-4">Custom Icons</a></span></li><li><span><a href="#Utility-Functions" data-toc-modified-id="Utility-Functions-5">Utility Functions</a></span></li></ul></div>

# 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
import pandas as pd

import msticpy.sectools as sectools
from msticpy.nbtools import *
from msticpy.nbtools.entityschema import IpAddress, GeoLocation
from msticpy.nbtools.foliummap import FoliumMap


# ## FoliumMap class
# 
# ```
# FoliumMap(
#     title: str = 'layer1',
#     zoom_start: float = 2.5,
#     tiles=None,
#     width: str = '100%',
#     height: str = '100%',
#     location: list = None,
# )
# Wrapper class for Folium/Leaflet mapping.
# 
# Parameters
# ----------
# title : str, optional
#     Name of the layer (the default is 'layer1')
# zoom_start : int, optional
#     The zoom level of the map (the default is 7)
# tiles : [type], optional
#     Custom set of tiles or tile URL (the default is None)
# width : str, optional
#     Map display width (the default is '100%')
# height : str, optional
#     Map display height (the default is '100%')
# location : list, optional
#     Location to center map on
# 
# Attributes
# ----------
# folium_map : folium.Map
# ```

# In[2]:


folium_map = FoliumMap(width="50%", height="50%", location=(47.5982328,-122.331), zoom_start=14)
folium_map


# The underlying folium map object is accessible as the `folium_map` attribute

# In[3]:


type(folium_map.folium_map)


# ## Adding IP Entities to the map
# 
# ```
# fol_map.add_ip_cluster(
#     ip_entities: Iterable[msticpy.nbtools.entityschema.IpAddress],
#     **kwargs,
# )
# 
# Add a collection of IP Entities to the map.
# 
# Parameters
# ----------
# ip_entities : Iterable[IpAddress]
#     a iterable of IpAddress Entities
# 
# Other Parameters
# ----------------
#     kwargs: icon properties to use for displaying this cluster
# ```

# In[4]:


import pickle
with open(b"data/ip_entities.pkl", "rb") as fh:
    ip_entities = pickle.load(fh)
ip_entities = [ip for ip in ip_entities if ip.Location and ip.Location.Latitude]

folium_map = FoliumMap(zoom_start=9)
folium_map.add_ip_cluster(ip_entities=ip_entities, color='orange')
folium_map.center_map()
folium_map


# In[13]:


# Read in some data
geo_loc_df = pd.read_csv("data/ip_locs.csv", index_col=0)
geo_loc_df.head()


# In[14]:


# Create IP and GeoLocation Entities from the dataframe
def create_ip_entity(row):
    ip_ent = IpAddress(Address=row["AllExtIPs"])
    geo_loc = create_geo_entity(row)
    ip_ent.Location = geo_loc
    return ip_ent

def create_geo_entity(row):
    # get subset of fields for GeoLocation
    loc_props = row[["CountryCode", "CountryName","State", "City", "Longitude", "Latitude"]]
    geo_loc = GeoLocation(**loc_props.to_dict())
    return geo_loc

geo_locs = list(geo_loc_df.apply(create_geo_entity, axis=1).values)
ip_ents = list(geo_loc_df.apply(create_ip_entity, axis=1).values)
ip_ents[:5]


# In[15]:


geo_loc_df.apply(lambda x: (x.Latitude, x.Longitude), axis=1).values


# ### Plot IPAddress entities with location data

# In[16]:


fmap_ips = FoliumMap()
fmap_ips.add_ip_cluster(ip_entities=ip_ents[:20], color='blue')
fmap_ips.center_map()
fmap_ips


# ### Use different colors and icons

# In[17]:


fmap_ips.add_ip_cluster(ip_entities=ip_ents[30:40], color='red', icon="flash")
fmap_ips.center_map()
fmap_ips


# ## Custom Icons
# 
# By default folium uses the information icon (i).
# Icons can be taken from the default Bootstrap set. See the default list here [glyphicons](https://www.w3schools.com/icons/bootstrap_icons_glyphicons.asp)
# 
# Alternatively you can use icons from the [Font Awesome collection](https://fontawesome.com/icons?d=gallery)
# by adding prefx="fa" and icon="icon_name" to the call to add_ip_cluster or add_geo_cluster.

# In[18]:


fmap_ips.add_geoloc_cluster(geo_locations=geo_locs[40:50], color='darkblue', icon="desktop", prefix="fa")
fmap_ips.center_map()
fmap_ips


# ## Utility Functions

# ### Calculate center point of entity locations

# In[19]:


from msticpy.nbtools.foliummap import get_map_center, get_center_ip_entities, get_center_geo_locs

print(get_center_geo_locs(geo_locs))
print(get_center_geo_locs(geo_locs, mode="mean"))

# get_map_center Will accept iterable of any entity type that is either
# an IpAddress entity or an entity that has properties of type IpAddress
print(get_map_center(ip_ents[30:40]))
print(get_map_center(ip_ents[:20]))

print(get_center_ip_entities(ip_ents[:20])) 


# ### Calculate distance between entity locations

# In[27]:


from msticpy.sectools.geoip import entity_distance
print("Distance between")
print(ip_ents[0], ip_ents[1])
print("\n", entity_distance(ip_ents[0], ip_ents[1]), "km")

print("Distance between")
print(ip_ents[0], ip_ents[13])
print("\n", entity_distance(ip_ents[0], ip_ents[13]), "km")