content: {-
      Copyright 2016 Markus Ongyerth
  
      This file is part of Monky.
  
      Monky is free software: you can redistribute it and/or modify
      it under the terms of the GNU Lesser General Public License as published by
      the Free Software Foundation, either version 3 of the License, or
      (at your option) any later version.
  
      Monky is distributed in the hope that it will be useful,
      but WITHOUT ANY WARRANTY; without even the implied warranty of
      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
      GNU Lesser General Public License for more details.
  
      You should have received a copy of the GNU Lesser General Public License
      along with Monky.  If not, see <http://www.gnu.org/licenses/>.
  -}
  {-# LANGUAGE TupleSections #-}
  {-|
  Module      : Monky.IP
  Description : IP information collection module
  Maintainer  : ongy
  Stability   : experimental
  Portability : Linux
  
  Some of this should (and will) move to the netlink package at some point in time.
  
  This mostly exists, because both i3-status and conky have
  modules like this.
  IPs have become complicated and I don't really know how
  to sanely display IP in formation for an interface.
  Feel free to create issues/contact me on IRC if you
  want to use this and need a better interface (maybe network length information?)
  -}
  module Monky.IP
    ( IPHandle
    , getSocket
    , getAddresses
  
    , AddressFamily(..)
    , IP(..)
    , getRawFd
    , subscribeToEvents
    , handleNext
    )
  where
  
  import System.IO (hPutStrLn, stderr)
  import Data.Bits ((.|.))
  import Data.Maybe (mapMaybe)
  import Data.Word (Word32)
  
  import System.Posix.Types (Fd)
  
  import System.Linux.Netlink
  import System.Linux.Netlink.Constants hiding (AddressFamily)
  import System.Linux.Netlink.Route
  
  import qualified Data.ByteString.Char8 as BSC (pack)
  import qualified Data.ByteString as BS (singleton, append)
  
  import qualified Data.Map as M
  
  import Monky.IP.Raw
  
  -- |Handle to access information about the IPs on the system
  data IPHandle = IPHandle NetlinkSocket AddressFamily Word32
  
  -- This is RTNLGRP_IPV4_IFADDR, which is not currently exported by the netlink library
  cRTNLGRP_IPV4_IFADDR :: Num a => a
  cRTNLGRP_IPV4_IFADDR = 5
  
  cRTNLGRP_IPV6_IFADDR :: Num a => a
  cRTNLGRP_IPV6_IFADDR = 9
  
  linkQuery :: String -> RoutePacket
  linkQuery name =
    let flags = fNLM_F_REQUEST -- .|. fNLM_F_MATCH .|. fNLM_F_ROOT
        header = Header eRTM_GETLINK flags 0 0
        msg = NLinkMsg 0 0 0
        attrs = M.fromList [(eIFLA_IFNAME, BSC.pack name `BS.append` BS.singleton 0)]
        in
      Packet header msg attrs
  
  
  getInterfaceID :: NetlinkSocket -> String -> IO Word32
  getInterfaceID sock name = do
    interfaces <- query sock $ linkQuery name
    let ids = map (\(Packet _ (NLinkMsg _ index _) _) -> index) interfaces
    return $ head ids
  
  -- |Get an IPHandle to gather information about IPs on the system.
  getSocket
    :: String -- ^The interface to monitor
    -> AddressFamily -- ^The AddressFamily to use
    -> IO IPHandle
  getSocket name fam = do
    sock <- makeSocket
    iid <- getInterfaceID sock name
    return $ IPHandle sock fam iid
  
  -- I'm not quite sure why, but apparently RTM_GETADDR only has a version for dumps, not
  -- normal requests, so we have to filter the output either way.
  addressQuery :: AddressFamily -> RoutePacket
  addressQuery fam =
    let flags = fNLM_F_REQUEST .|. fNLM_F_MATCH .|. fNLM_F_ROOT
        header = Header eRTM_GETADDR flags 0 0
        msg = NAddrMsg (familyToNum fam) 0 0 0 0
        attrs = M.empty
        in
      Packet header msg attrs
  
  -- |Get all addresses on the system (filtered by interface and type, see 'getSocket')
  getAddresses :: IPHandle -> IO [IP]
  getAddresses (IPHandle sock fam iid) = do
    packs <- query sock $ addressQuery fam
    let matching =  filter (\(Packet _ (NAddrMsg _ _ _ _ aid) _) -> aid == iid) packs
    let addrs = mapMaybe (M.lookup eIFLA_ADDRESS . packetAttributes) matching
    return $ map (ipFromBS) addrs
  
  -- |Get the raw 'Fd' used by this handle. This can be used to block in RTS controlled manner
  getRawFd :: IPHandle -> Fd
  getRawFd (IPHandle sock _ _) = getNetlinkFd sock
  
  -- |Subscribe to the multicast group(s) for this handle. This has to be called ONCE before 'handleNext' can be used.
  subscribeToEvents :: IPHandle -> IO ()
  subscribeToEvents (IPHandle sock fam _) =
    let grps = case fam of
                 AF_UNSPEC -> [cRTNLGRP_IPV6_IFADDR, cRTNLGRP_IPV4_IFADDR]
                 AF_INET -> [cRTNLGRP_IPV4_IFADDR]
                 AF_INET6 -> [cRTNLGRP_IPV6_IFADDR]
                 in
      mapM_ (joinMulticastGroup sock) grps
  
  
  -- Handle -> add -> remove -> done
  {- |Handle the next event from Netlink
  
  This will handle the next event reported by the handle.
  Before this ever does anything 'subscribeToEvents' has to be called ONCE!
  This function can then be called as often as needed and should be called
  in a main loop.
  
  The functions passed may be called multiple times, and sometimes both may be called.
  This is a limitation by the netlink api, that cannot be avoided without user space buffers.
  -}
  handleNext
    :: IPHandle -- ^The handle to use
    -> (IP -> IO ()) -- ^Function to call with new ip addresses
    -> (IP -> IO ()) -- ^Function to call with removed addresses
    -> IO ()
  handleNext (IPHandle sock _ iid) add remove = do
    evt <- recvOne sock
    let matching = filter (\(Packet _ (NAddrMsg _ _ _ _ aid) _) -> aid == iid) evt
    let addrs = mapMaybe transform matching
    mapM_ handle addrs
    where handle (cmd, addr)
            | cmd == eRTM_NEWADDR = add addr
            | cmd == eRTM_DELADDR = remove addr
            | otherwise = hPutStrLn stderr ("Got unexpeced message while handling IPevents: " ++ show cmd)
          transform :: RoutePacket -> Maybe (MessageType, IP)
          transform (Packet (Header cmd _ _ _) _ attrs) =
            fmap ((cmd,) . ipFromBS) . M.lookup eIFLA_ADDRESS $ attrs
          transform err = error ("Something went wrong while handling IP events: " ++ show err)
  