-- | This module provides facilities for easily creating Netkit (<http://wiki.netkit.org/index.php/Main_Page>) labs  
-- to experiment with OpenFlow switches and controllers. 
-- This module provides a notation for describing simple 
-- OpenFlow-based network topologies, and provides a command
-- that generates a NetKit lab that can be used to simulate the
-- network.
--
-- To use this module, describe the configuration of your test network, including
-- hosts, switches and their interconnectivity using the functions in this module (see the example below), 
-- and then run one of the @makeLab@ commands to generate the Netkit files needed to run the lab.
-- You can then move into the directory containing the Netkit lab and start the lab using Netkit commands 
-- (e.g. @lstart@).
-- The generated lab will include hosts and switches, but not a controller. 
-- The lab will setup a TAP interface with subnet 10.0.0.0/8 from the switch virtual machines to the host on which you started the lab, 
-- and the switches will attempt to contact a controller with IP address 10.0.0.1 over that TAP interface, at the server port mentioned in the lab.
-- You can then control the switches by starting a controller on the host at the specified port. See the example below for more details.
-- 
-- The generated lab is designed to work with a customized netkit file system that 
-- has OpenFlow software in a particular location. Instructions for obtaining this file
-- system are here <http://haskell.cs.yale.edu/?page_id=383>; see the last instruction in the section on Installing on Your Own Machine. 
-- The generator (i.e. @makeLab@) must know the location of this file in order to generate the Netkit lab files. 
-- The default options assume the files are in the user's @~/.nettle@ directory, but they can be placed in other locations as well. 
-- If they are in another location, then the Lab options must be set appropriately. 
module Nettle.Netkit.LabUtil 
    (
     -- * Lab descriptions
      LabConfig (..)
    , Options (..)
    , OpenFlowVersion (..)
    , ControllerTCPPort
    , Switch (..)
    , Host (..)
    , Port (..)
    , Interface (..)
    , SwitchPort (..)
    , switch
    , host
    , port
    , (#)
    , (<-->)
    , (@@)

    -- * Generating labs
    , makeLab
    , makeLabWithDefaults
    , getDefaultOptions
    -- * Example
    -- $doc-example 

    -- * Limitations
    -- $doc-limitations 
    )
    where

import Nettle.OpenFlow.Switch (SwitchID)
import Nettle.OpenFlow.Port hiding (Port)
import Nettle.IPv4.IPAddress
import Data.List (nub, elemIndex, find, union, sort, (\\), minimumBy)
import System.Directory (createDirectory, createDirectoryIfMissing, getHomeDirectory, canonicalizePath)
import Control.Monad.State
import Data.Maybe (fromJust)
import Data.Bits
import Data.Map (Map)
import qualified Data.Map as Map
import System.FilePath
import System.Posix.Files
import System.Posix.Types
import Data.Set (Set)
import qualified Data.Set as Set
import Nettle.Netkit.UnionFind

newtype Switch = Switch SwitchID deriving (Show, Eq, Ord)

data Host = Host { hostID :: Int } deriving (Show, Eq, Ord)

data Interface = Interface { interfaceID :: Int, 
                             interfaceIPAddress :: IPAddress } 
                 deriving (Show,Eq,Ord)

newtype Port   = Port PortID deriving (Show, Eq, Ord)

type SwitchPort = (Switch,Port)
type ControllerTCPPort = Int

switch :: SwitchID -> Switch
switch n = Switch n

host :: Int -> Host
host n = Host n 

port :: PortID -> Port
port n = Port n

data LabConfig = 
    LabConfig { switches             :: [Switch]                        -- ^ Switches for this lab
              , controllerServerPort :: ControllerTCPPort               -- ^ The TCP port number at which the controller will listen for switch connections
              , hosts                :: [((Host,Interface),SwitchPort)] -- ^ A description of where the host interfaces are attached to switches in the network
              , links                :: [(SwitchPort,SwitchPort)]       -- ^ A description of how switches are connected
              } deriving (Show,Eq)

-- | A datatype for specifying lab options, including the paths to the kernel and filesystem
-- used by netkit machines, as well as the OpenFlow version of the reference switch software.
data Options = Options { pathToKernel, pathToFileSystem :: Maybe String
                       , openFlowVersion :: OpenFlowVersion } 
             deriving (Show,Eq)

-- | An enumerated data type representing OpenFlow versions supported by 
-- this module.
data OpenFlowVersion = Ver0_9_0 
                     | Ver1_0_0 
                     deriving (Eq,Show,Ord,Enum)

getDefaultOptions :: OpenFlowVersion -> IO Options
getDefaultOptions version = do 
  home <- getHomeDirectory
  return (Options { pathToKernel     = Nothing, 
                    pathToFileSystem = Just (fsdir home), 
                    openFlowVersion  = version      })

  where fsdir base = base </> 
                     ".nettle" </> 
                     "netkit" </> 
                     "netkit-filesystem-i386-F5.1_2010-11-19" </>
                     "netkit-fs"

infix 8 #
infix 7 @@, <-->

-- | Denotes a SwitchPort, i.e. a port on a switch.
(#) :: Switch -> Port -> SwitchPort
(#) s p = (s,p)

-- | Denotes where a host is attached to the network of switches.
(@@) :: (Host, Interface) -> SwitchPort -> ((Host,Interface), SwitchPort)
h @@ sp = (h,sp)

-- | Denotes a link (switch-to-switch) connection.          
(<-->) :: SwitchPort -> SwitchPort -> (SwitchPort, SwitchPort)
(<-->) a b = (a,b)


-- | @makeLabWithDefaults path options lab@ is a command that writes the files and
-- directories needed to run a Netkit lab that implements the description provided by @lab@. 
-- It writes the files to directory @path@, and the options are specified by @options@.
makeLab :: FilePath -> Options -> LabConfig -> IO ()
makeLab path options lab = 
    do path'  <- canonicalizePath path
       layout <- lab2NetkitLayout path' options lab
       makeNetkitLayout path' layout

-- | @makeLabWithDefaults path version lab@ is a command that writes the files and
-- directories needed to run a Netkit lab that implements the description provided by @lab@. 
-- It writes the files to directory @path@, and the switches will run the OpenFlow reference
-- switch software for OpenFlow version @version@.
makeLabWithDefaults :: FilePath -> OpenFlowVersion -> LabConfig -> IO ()
makeLabWithDefaults path version lab = 
    do options <- getDefaultOptions version
       makeLab path options lab

-- Supporting functions and types...

-- A value of type NetkitLayout represents the 
-- directories and files that need to be created
-- for a lab.
data NetkitLayout = 
    NetkitLab { dirs :: [FilePath],
                files :: [(FilePath, String, Maybe FileMode)] }
    deriving (Show,Read,Eq)

-- Makes the directories and files described in a NetkitLayout value.
makeNetkitLayout :: FilePath -> NetkitLayout -> IO ()
makeNetkitLayout path labLayout = do 
  createDirectoryIfMissing True path
  sequence_ [ createDirectory d |     d <- dirs labLayout ]
  sequence_ [ createDirectoryIfMissing True parent >> 
              writeFile f s >> 
              case m of {Nothing -> return (); Just mode -> setFileMode f mode} 
              | (f,s,m) <- files labLayout, let (parent,fname) = splitFileName f ]


-- Returns a list of all switches explicitly mentioned
-- in a configuration.
switches' :: LabConfig -> [Switch]
switches' c = nub ( [ sw | (_,(sw,_)) <- hosts c ] ++ 
                    [ sw | ((sw,_),_) <- links c ] ++ 
                    [ sw | (_,(sw,_)) <- links c ] ++ 
                    switches c)

uniqueHosts :: LabConfig -> [Host]
uniqueHosts = nub . map (fst . fst) . hosts

counts :: Eq a => [a] -> [(a,Int)]
counts xs = [ (x, length (filter (==x) xs)) | x <- nub xs ]


lab2NetkitLayout :: FilePath -> Options -> LabConfig -> IO NetkitLayout
lab2NetkitLayout labdir options conf 
    | not (null occMoreThanOnce) = error ("The following host IDs are used more than once for different hosts: " ++ show occMoreThanOnce)
    | otherwise = 
        do fsPath  <- maybe (return Nothing) ( return.Just <=< canonicalizePath ) (pathToFileSystem options)
           kerPath <- maybe (return Nothing) ( return.Just <=< canonicalizePath ) (pathToKernel     options)
           return $ NetkitLab dirs (map (\(fn,cont, m) -> (labdir </> fn, cont, m)) (files fsPath kerPath))
    where occMoreThanOnce      = [ hid | (hid,cnt) <- counts [ hostID h | h <- uniqueHosts conf], cnt > 1]
          dirs                 = [ labdir </> d | d <- (dirsForSwitches ++ dirsForHosts) ]
          dirsForSwitches      = [ switchName s | s <- sws ]
          dirsForHosts         = [ hostName h   | h <- uniqueHosts conf ]
          files fsPath kerPath = [ labConfFile fsPath kerPath]
                                 ++ map (hostStartupFile conf . fst . fst) hostConns 
                                 ++ map switchStartupFile switchControlConfig
                                 ++ concatMap (switchScripts conf options switchControlConfig) sws
          labConfFile fsPath kerPath = ("lab.conf", makeLabConfContents conf fsPath kerPath, Nothing)
          switchControlConfig  = configToSwitchControlConfigs conf
          sws                  = switches' conf
          hostConns            = hosts conf
          swConns              = links conf
          

{- Scripts that should be available at the switches. -}
switchScripts conf options switchControlConfig sw = 
    [ 
     (switchName sw ++ "/startSwitch.sh"          , startSwitchContents,                                Just mode), 
     (switchName sw ++ "/startProtocol.sh"        , startProtocolContents,                              Just mode), 
     (switchName sw ++ "/etc/modprobe.d/blacklist", blacklistContents,                                  Just mode), 
     (switchName sw ++ "/root/.profile"           , "export PATH=$PATH:" ++ ofpBinDir version, Just mode)
    ]
    where startSwitchContents        = makeScriptFile (ofdatapathCommand version sw [ "eth" ++ show p | Port p <- ports sw ])
          startProtocolContents      = makeScriptFile (ofprotocolCommand version controllerAddressForSwitch controllerTCPPort)
          blacklistContents          = unlines ["blacklist net-pf-10", "blacklist ipv6"]
          ports sw                   = sort $ fillMissingPorts $ connectedPorts conf sw
          controllerAddressForSwitch = 
              let (_,_,_,addr) = fromJust $ find (\(sw',_,_,_) -> sw==sw') switchControlConfig
              in addr
          mode              = ownerModes `unionFileModes` groupModes
          controllerTCPPort = controllerServerPort conf
          sws               = switches' conf
          hostConns         = hosts conf
          swConns           = links conf
          version           = openFlowVersion options 

fillMissingPorts :: [Port] -> [Port]
fillMissingPorts ps = ps `union` [ Port i | i <- [1..maximum pnums]]
    where pnums = [ pnum | Port pnum <- ps ]

makeScriptFile cmd = "#!/bin/bash\n\n" ++ cmd ++ "\n"

ofpBinDir :: OpenFlowVersion -> FilePath
ofpBinDir Ver0_9_0 = "/root/local/openflow-0.9.0/bin"
ofpBinDir Ver1_0_0 = "/root/local/openflow-1.0.0/bin"

ofdatapathCommand version (Switch dpid) ifaces 
    = ofdatapathPath version ++ 
      " -i " ++ concat (sepBy "," ifaces) ++ 
      " punix:ofconnfile -d " ++ showDataPathID dpid ++ 
      " --no-local-port " ++ 
      if version == Ver1_0_0 then "--no-slicing" else "" ++ ";" 
ofdatapathPath version = ofpBinDir version </> "ofdatapath"

ofprotocolCommand version caddr cport = 
    ofprotocolPath version ++ " unix:ofconnfile tcp:" ++ showOctets caddr ++ ":" ++ show cport ++ " -v -Fclosed"
ofprotocolPath version = ofpBinDir version </> "ofprotocol"

showDataPathID :: SwitchID -> String
showDataPathID n = replicate (totalLen - l) '0' ++ digs
    where l        = length digs
          totalLen = 12
          digs     = hexDigits (fromIntegral n)

hexDigits :: Int -> [Char]
hexDigits = reverse . map (intToHexDigit . (`mod` base)) . takeWhile (>0) . iterate (`div` base)
        where base = 16

intToHexDigit :: Int -> Char
intToHexDigit n
    | n >= 0 && n < 10 = head $ show n
    | n == 10          = 'a'
    | n == 11          = 'b'
    | n == 12          = 'c'
    | n == 13          = 'd'
    | n == 14          = 'e'
    | n == 15          = 'f'

sepBy _ [] = []
sepBy _ [x] = [x]
sepBy s (x:xs) = x : s : sepBy s xs

{- Start-up file for switches -}
switchStartupFile :: (Switch, Port, IPAddress, IPAddress) -> (FilePath, String, Maybe FileMode)
switchStartupFile (sw, Port nic,_,_) = (switchName sw ++ ".startup", contents, Nothing)
    where contents = 
              line (defaultRouteCommand controlInterface) ++
              startSwitchCommands ++ 
              startProtocolCommands           
          controlInterface = "eth" ++ show nic
          startSwitchCommands = line "chmod +x /startSwitch.sh" ++ 
                                line "screen -S startSwitch -d -m /startSwitch.sh "
          startProtocolCommands = line "chmod +x /startProtocol.sh" ++ 
                                  line "screen -S startProtocol -d -m /startProtocol.sh "

line :: String -> String
line s = s ++ "\n"

{- Start-up file for hosts -}
-- hostStartupFile :: LabConfig -> Host -> 
hostStartupFile lab h = (hostName h ++ ".startup", contents, Nothing)
    where contents = 
              concat [ line (ifConfigCommand' (iName i) (showOctets (interfaceIPAddress i)) "255.255.255.0" ) | i <- ints ]  ++
              line (defaultRouteCommand $ iName $ minimumBy (\a b -> compare (interfaceID a) (interfaceID b)) ints)
          h'      = fromIntegral (hostID h)
          iName i = "eth" ++ show (interfaceID i)
          ints = hostInterfaces lab h
          

hostInterfaces :: LabConfig -> Host -> [Interface]
hostInterfaces lab host = [ i | ((host',i),_) <- hosts lab, host' == host]

defaultRouteCommand interfaceName = "route add default " ++ interfaceName

{-
--          interfaceName = "eth0"
-- line (ifConfigCommand interfaceName (showOctets (ipAddress 11 1 h' 5)) "255.255.255.0" (showOctets (ipAddress 11 1 h' 11)))

-}
ifConfigCommand ifaceName address mask bcastAddress = 
    "ifconfig " ++ ifaceName ++ " " ++ address ++ " netmask " ++ mask ++ " broadcast " ++ bcastAddress ++ " up"

ifConfigCommand' ifaceName address mask = 
    "ifconfig " ++ ifaceName ++ " " ++ address ++ " netmask " ++ mask ++ " up"


{- Function to make the lab.conf file -}

makeLabConfContents :: LabConfig -> Maybe FilePath -> Maybe FilePath -> String
makeLabConfContents conf fsPath kerPath = 
    output $ execState (labConfContents conf fsPath kerPath) initState
    where initState = ConfWriterState { output = "" }

type ConfWriter a = State ConfWriterState a
data ConfWriterState = ConfWriterState { output :: String } 
                     deriving (Show,Eq)

addLine :: String -> ConfWriter ()
addLine s = modify (\state -> state { output = output state ++ line s })

labConfContents :: LabConfig -> Maybe FilePath -> Maybe FilePath -> ConfWriter ()
labConfContents lab fsPath kerPath = do
  mapM_ addHostSwitchConnectionLines hostConns
  addSwitchInterfaces
  mapM_ addSwitchControllerLine (configToSwitchControlConfigs lab)
  addKernelAndFSOptions
  setMemoryParams

    where 

      setMemoryParams = 
          do sequence_ [ addLine (confLine (hostName h) "mem" "512") | h <- uniqueHosts lab ]

      addKernelAndFSOptions = 
        let f name = do maybe (return ()) (addLine . confLine name "model-fs") fsPath
                        maybe (return ()) (addLine . confLine name "kernel") kerPath
        in do sequence_ [ f (hostName h)    | h <- uniqueHosts lab ]
              sequence_ [ f (switchName sw) | sw <- sws ]

      addSwitchInterfaces = 
          sequence_ [ addLine (confLine (switchName sw) (show p) [label])
                      | ((sw, Port p), label) <- Map.assocs spLabels ]

      addHostSwitchConnectionLines ((h,i),(s, Port p)) = do 
        let label = spLabel s (Port p)
        addLine (confLine (hostName h) (show (interfaceID i)) [label])

      addSwitchControllerLine (sw, Port nic, nicIP, controlIP) = do 
        let Switch s = sw
        let switchIPAddress = showOctets nicIP
        let controlIPAddress = showOctets controlIP
        let c = "tap," ++ controlIPAddress ++ "," ++ switchIPAddress 
        addLine (confLine (switchName sw) (show nic) c)

      confLine machineName paramName value = 
          machineName ++ "[" ++ paramName ++ "]=" ++ value

      sws          = switches' lab
      hostConns    = hosts lab
      swConns      = links lab
      spLabels     = labelledParts lab
      spLabel sw p = Map.findWithDefault (error ("unknown (switch,port): " ++ show (sw,p))) (sw,p) spLabels


unConnectedInterfaces :: LabConfig -> [(Switch,Port)]
unConnectedInterfaces lab = 
    [ (sw,p) | sw <- switches' lab, p <- unConnectedPorts lab sw ]

allSwitchPorts :: LabConfig -> [(Switch,Port)]
allSwitchPorts lab = [ (sw,port) | sw <- switches' lab, port <- connectedPorts lab sw] ++ 
                     [ (sw,port) | sw <- switches' lab, port <- unConnectedPorts lab sw]
                     
labelledParts :: LabConfig -> Map (Switch, Port) Char 
labelledParts lab = Map.fromList keyVals 
  where startRel = links lab ++ [ (sif, sif) | (_, sif) <- hosts lab ]
        partCharPairs = zip (Set.elems (finestPartition startRel)) ['a','b'..'z'] 
        keyVals = [ (sp, c) | (spSet, c) <- partCharPairs, sp <- Set.elems spSet ]

type Label = String

connectedPorts :: LabConfig -> Switch -> [Port]
connectedPorts lab sw = nub (   [ p | (_,(sw',p))  <- hosts lab, sw==sw' ]
                             ++ [ p | ((sw',p), _) <- links lab, sw==sw' ]
                             ++ [ p | (_, (sw',p)) <- links lab, sw==sw' ])

unConnectedPorts :: LabConfig -> Switch -> [Port]
unConnectedPorts lab sw = [ Port i | i <- [1..maximum pnums] ] \\ ps
    where pnums = [ pnum | Port pnum <- ps ]
          ps    = connectedPorts lab sw

switchName :: Switch -> String
switchName (Switch s) = "switch" ++ show s

hostName :: Host -> String
hostName h = "host" ++ show (hostID h)

type SwitchControlConfig = [(Switch, Port, IPAddress, IPAddress)]

configToSwitchControlConfigs :: LabConfig -> SwitchControlConfig
configToSwitchControlConfigs conf = 
    [ (sw, 
       switchSidePort sw, 
       switchSideAddress (fromIntegral swnum), 
       controlSideAddress (fromIntegral swnum))
      | sw@(Switch swnum) <- sws ] 
    where switchSideAddress sw  = ipAddress 10 0 0 (1 + sw)
          controlSideAddress sw = ipAddress 10 0 0 1
          switchSidePort sw     = port $ maximum [ p | Port p <- portsUsed sw] + 1
          portsUsed sw          = [ p | (h,(sw',p)) <- hostConns, sw==sw' ] ++ 
                                  [ p | ((sw',p),_) <- swConns, sw==sw'   ] ++ 
                                  [ p | (_, (sw',p)) <- swConns, sw==sw'   ] 
    
          sws = switches' conf
          hostConns = hosts conf
          swConns = links conf



-- $doc-example 
-- Below is an example lab config, describing a network with 
-- 3 switches, each connected to the other two and 3 hosts (one per switch).
-- The switches will attempt have one interface associated with a TAP interface on the 
-- host system. This TAP interface has subnet 10.0.0.0/8 and the controller is assumed to 
-- be running at 10.0.0.1. In this example, the controller should be running at TCP 
-- port 2525, so that the switches will find it. 
-- 
-- > lab :: LabConfig
-- > lab = LabConfig { controllerServerPort = 2525, 
-- >                   switches             = [sw1, sw2, sw3], 
-- >                   hosts                = [ (host 1, Interface 0 (ipAddress 11 1 1 5)) @@ (sw1 # port 1), 
-- >                                            (host 2, Interface 0 (ipAddress 11 1 2 5)) @@ (sw2 # port 1), 
-- >                                            (host 3, Interface 0 (ipAddress 11 1 3 5)) @@ (sw3 # port 1) ], 
-- >                   links                = [sw1 # port 2 <--> sw2 # port 2, 
-- >                                           sw1 # port 3 <--> sw3 # port 2,
-- >                                           sw2 # port 3 <--> sw3 # port 3 ] 
-- >                               }
-- >    where sw1 = switch 1
-- >          sw2 = switch 2
-- >          sw3 = switch 3 
--
-- The lab can be generated and written to directory \/foo\/bar by running the following command (the target directory may need to be created before running this command):
--
-- > makeLabWithDefaults "/foo/bar" Ver1_0_0 lab
--
-- The switches will have two files (both in the root directory), @startProtocol.sh@, @startSwitch.sh@ 
-- which run the ofprotocol and ofdatapath programs, respectively, with the correct
-- parameters for the switch and controller. The switches will run these scripts at start up time.
-- These scripts are run with @screen@, and you can reattach to them using @screen@. To see
-- the screen session names, run @screen -list@ from one of the switch terminals.

-- $doc-limitations
-- The library currently has several limitations: 
--  * Each host interface is on a /24 subnet containing its IP address. This should be configurable.
--  * The controller is assumed to be at 10.0.0.1. This should be configurable.