agent

agent.go

@@ -9,8 +9,9 @@ import (
 	dat "loafle.com/overflow/data_sender_go"
 	evt "loafle.com/overflow/event_sender_go"
 	pol "loafle.com/overflow/long_poller_go"
-	"log"
+	log "github.com/cihub/seelog"
 	"sync"
+	"loafle.com/overflow/agent/bootstrap"
 )
 
 const (
@@ -21,7 +22,6 @@ const (
 var wg sync.WaitGroup
 
 func main() {
-	log.SetPrefix("Agent : ")
 	wg.Add(1)
 	go StartAgent()
 	wg.Wait()
@@ -33,15 +33,21 @@ type Agent struct {
 }
 
 func StartAgent() {
+	stopch := make(chan bool, 1)
+	bootstrap.HandleShell(stopch)
+	bootstrap.HandleSignal(stopch)
+
 	agt := Agent{}
 	agt.startAgent()
+
+	agt.handleAgentStop(stopch)
 }
 
 func (agt *Agent) startAgent() {
 	//1. cfg start()
 	cfgStarted := make(chan *config_manager.GlobalConfig, 1)
 	if err := cfg.Start(cfgStarted, PATH); err != nil {
-		log.Println(err)
+		log.Error(err)
 		return
 	}
 	globalConf := <-cfgStarted
@@ -49,7 +55,7 @@ func (agt *Agent) startAgent() {
 	//2. evt start()
 	evtStarted := make(chan bool, 1)
 	if err := evt.Start(evtStarted); err != nil {
-		log.Println(err)
+		log.Error(err)
 		return
 	}
 	<-evtStarted
@@ -62,7 +68,7 @@ func (agt *Agent) startAgent() {
 	//4. pol start()
 	polStarted := make(chan bool, 1)
 	if err := pol.Start(polStarted, globalConf); err != nil {
-		log.Println(err)
+		log.Error(err)
 		return
 	}
 	<-polStarted
@@ -75,7 +81,7 @@ func (agt *Agent) startAgent() {
 	//5. scm start()
 	scmStarted := make(chan config_manager.ConfigManager, 1)
 	if err := cfg.StartSensorConfig(scmStarted, "temp"); err != nil {
-		log.Println(err)
+		log.Error(err)
 		return
 	}
 	agt.cm = <-scmStarted
@@ -99,6 +105,18 @@ func (agt *Agent) startAgent() {
 	go agt.waitAgentTask()
 }
 
+func (agt *Agent) handleAgentStop(ch chan bool) {
+	go func() {
+		if <-ch {
+			agt.stopAgent()
+			log.Info("Agent has stopped.")
+			log.Flush()
+			close(ch)
+			wg.Done()
+		}
+	}()
+}
+
 func (agt *Agent) waitAgentTask() {
 	for {
 		d := <-agt.taskCh
@@ -175,18 +193,24 @@ func (agt *Agent) stopAgent() {
 	evt.AddEventData(msg.Event{msg.EVT_TYPE_NONE, msg.AGT_STOP, nil})
 	evt.Stop()
 
-	wg.Done()
 }
 
 func (agt *Agent) processStartSensor(task msg.AgentTask) {
 	id := task.Params["sensorId"]
-	col.StartSensor(id)
+	if err := col.StartSensor(id); err != nil {
+		log.Error(err)
+		return
+	}
+
 	evt.AddEventData(msg.Event{msg.EVT_TYPE_TASK, task.Params["id"], nil})
 }
 
 func (agt *Agent) processStopSensor(task msg.AgentTask) {
 	id := task.Params["sensorId"]
-	col.StopSensor(id)
+	if err := col.StopSensor(id); err != nil {
+		log.Error(err)
+		return
+	}
 	evt.AddEventData(msg.Event{msg.EVT_TYPE_TASK, task.Params["id"], nil})
 }
 
@@ -206,7 +230,7 @@ func (agt *Agent) processRemoveSensor(task msg.AgentTask) {
 	id := task.Params["sensorId"]
 	err := col.RemSensor(id)
 	if err != nil {
-		log.Println(err)
+		log.Error(err)
 		return
 	}
 	crm.RemoveSensor(id)
@@ -218,13 +242,13 @@ func (agt *Agent) processUpdateSensor(task msg.AgentTask) {
 	path := task.Params["path"]
 	id := task.Params["id"]
 	if err := col.StopSensor(id); err != nil {
-		log.Println(err)
+		log.Error(err)
 		return
 	}
 	cfg.UpdateSensor(path)
 	crm.UpdateSensor(id)
 	if err := col.UpdateSensor(id); err != nil {
-		log.Println(err)
+		log.Error(err)
 		return
 	}
 	evt.AddEventData(msg.Event{msg.EVT_TYPE_TASK, task.Params["id"], nil})
@@ -251,7 +275,7 @@ func (agt *Agent) processUpdatePolInterval(task msg.AgentTask) {
 	interval := task.Params["interval"]
 	err := pol.UpdateInterval(interval)
 	if err != nil {
-		log.Println(err)
+		log.Error(err)
 		return
 	}
 	evt.AddEventData(msg.Event{msg.EVT_TYPE_TASK, task.Params["id"], nil})

bootstrap/shell.go

@@ -0,0 +1,77 @@
+package bootstrap
+
+import (
+	log "github.com/cihub/seelog"
+	"net"
+	"os"
+	"strings"
+)
+
+const FILE_SERVER string = "/var/run/of_server"
+
+var fd net.Conn
+
+func HandleShell(stop chan bool) {
+	go func() {
+		//os.Stdout.Close()
+		//os.Stderr.Close()
+		//os.Stdin.Close()
+
+		if fi, _ := os.Stat(FILE_SERVER); fi != nil {
+			os.Remove(FILE_SERVER)
+		}
+
+		l, err := net.ListenUnix("unix", &net.UnixAddr{Name: FILE_SERVER, Net: "unix"})
+		if err != nil {
+			log.Error(err.Error())
+			return
+		}
+		defer func() {
+			l.Close()
+			os.Remove(FILE_SERVER)
+		}()
+
+		for {
+			fd, err = l.Accept()
+			if err != nil {
+				log.Error(err)
+				return
+			}
+
+			buf := make([]byte, 1024)
+			nr, err := fd.Read(buf)
+			if err != nil {
+				return
+			}
+
+			data := string(buf[0:nr])
+			switch strings.ToUpper(data) {
+			case "STOP":
+				if fi, _ := os.Stat(FILE_SERVER); fi != nil {
+					os.Remove(FILE_SERVER)
+				}
+				stop <- true
+			case "STATUS":
+				fd.Write(status())
+			default:
+			}
+
+			fd.Close()
+		}
+	}()
+}
+
+func StopHandleShell() {
+	log.Info("Shell Handler stopped.")
+
+	if fi, _ := os.Stat(FILE_SERVER); fi != nil {
+		os.Remove(FILE_SERVER)
+	}
+	if fd != nil {
+		fd.Close()
+	}
+}
+
+func status() []byte {
+	return []byte("STATUS OK\n")
+}

bootstrap/signal.go

@@ -0,0 +1,53 @@
+package bootstrap
+
+import (
+	log "github.com/cihub/seelog"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+func HandleSignal(stop chan bool) {
+	sigs := make(chan os.Signal, 1)
+
+	signal.Notify(sigs,
+		os.Kill,
+		os.Interrupt,
+		syscall.SIGKILL,
+		syscall.SIGSTOP,
+		syscall.SIGHUP,
+		syscall.SIGINT,
+		syscall.SIGTERM,
+		syscall.SIGQUIT,
+	)
+
+	go func() {
+		for {
+			s := <-sigs
+			switch s {
+			case os.Kill, os.Interrupt, syscall.SIGSTOP, syscall.SIGKILL, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT:
+				log.Infof("Signal received. [%s]", s)
+				stop <- true
+			default:
+				log.Infof("Signal received. [%s]", s)
+				stop <- true
+			}
+		}
+	}()
+}
+
+/*
+ 1) SIGHUP	 2) SIGINT	 3) SIGQUIT	 4) SIGILL	 5) SIGTRAP
+ 6) SIGABRT	 7) SIGBUS	 8) SIGFPE	 9) SIGKILL	10) SIGUSR1
+11) SIGSEGV	12) SIGUSR2	13) SIGPIPE	14) SIGALRM	15) SIGTERM
+16) SIGSTKFLT	17) SIGCHLD	18) SIGCONT	19) SIGSTOP	20) SIGTSTP
+21) SIGTTIN	22) SIGTTOU	23) SIGURG	24) SIGXCPU	25) SIGXFSZ
+26) SIGVTALRM	27) SIGPROF	28) SIGWINCH	29) SIGIO	30) SIGPWR
+31) SIGSYS	34) SIGRTMIN	35) SIGRTMIN+1	36) SIGRTMIN+2	37) SIGRTMIN+3
+38) SIGRTMIN+4	39) SIGRTMIN+5	40) SIGRTMIN+6	41) SIGRTMIN+7	42) SIGRTMIN+8
+43) SIGRTMIN+9	44) SIGRTMIN+10	45) SIGRTMIN+11	46) SIGRTMIN+12	47) SIGRTMIN+13
+48) SIGRTMIN+14	49) SIGRTMIN+15	50) SIGRTMAX-14	51) SIGRTMAX-13	52) SIGRTMAX-12
+53) SIGRTMAX-11	54) SIGRTMAX-10	55) SIGRTMAX-9	56) SIGRTMAX-8	57) SIGRTMAX-7
+58) SIGRTMAX-6	59) SIGRTMAX-5	60) SIGRTMAX-4	61) SIGRTMAX-3	62) SIGRTMAX-2
+63) SIGRTMAX-1	64) SIGRTMAX
+*/

script/ubuntu/of_agent

@@ -0,0 +1,71 @@
+### BEGIN INIT INFO
+# Provides: overFlow
+# Required-Start: $network
+# Required-Stop: $network
+# Default-Start: 2 3 4 5
+# Default-Stop: 0 1 6
+# Short-Description: overFlow agent
+# Description: overFlow Agent
+### END INIT INFO
+
+PROG="ofcollector"
+PROG_PATH="/home/insanity/develop/overflow/overflow.collector/bin/"
+PID_PATH="/var/run/"
+FILE_SERVER=/var/run/of_server
+
+start() {
+    if [ -e "$PID_PATH/$PROG.pid" ]; then
+        ## Program is running, exit with error.
+        echo "$PROG is currently running!" 1>&2
+        exit 1
+    else
+        $PROG_PATH/$PROG 2>&1 > /dev/null &
+        echo "$PROG started"
+        touch "$PID_PATH/$PROG.pid"
+    fi
+}
+
+stop() {
+    sudo echo -n 'STOP' | netcat -U $FILE_SERVER
+
+    if [ -e "$PID_PATH/$PROG.pid" ]; then
+        ## Program is running, so stop it
+        #killall $PROG
+        rm "$PID_PATH/$PROG.pid"
+        echo "$PROG stopped"
+    else
+        ## Program is not running, exit with error.
+        echo "Error! $PROG not started!" 1>&2
+        exit 1
+    fi
+}
+
+status() {
+    sudo echo -n $1 | netcat -U $FILE_SERVER
+}
+
+
+if [ "$(id -u)" != "0" ]; then
+    echo "This script must be run as root" 1>&2
+    exit 1
+fi
+
+case "$1" in
+    start)
+        start
+    ;;
+    stop)
+        stop
+    ;;
+    reload|restart|force-reload)
+        stop
+        start
+    ;;
+    status)
+        status 'STATUS'
+   ;;
+    **)
+        echo "Usage: $0 {start|stop|restart|status}" 1>&2
+        exit 1
+    ;;
+esac