From 37d938da17db2b7c19c5d4f582d35611f2b4113a Mon Sep 17 00:00:00 2001
From: Jonathan Sims <jonathan.625266@earthlink.net>
Date: Sat, 15 Feb 2020 12:25:29 -0500
Subject: [PATCH] Add option to run a script when client connects or
 disconnects (Linux only for now)

---
 README.md |  3 +++
 client.go | 20 ++++++++++++++++++++
 main.go   | 27 +++++++++++++++++++++++++--
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 43f31e63..63e26b6b 100644
--- a/README.md
+++ b/README.md
@@ -14,6 +14,7 @@ Features:
 * Each stream can have multiple video and audio tracks
 * Supports the RTP/RTCP streaming protocol
 * Optional authentication mechanism for publishers
+* Optionally run a script on client connect or disconnect (currently Linux only)
 * Compatible with Linux and Windows, does not require any dependency or interpreter, it's a single executable
 
 ## Installation
@@ -73,6 +74,8 @@ Flags:
   --rtp-port=8000        port of the RTP UDP listener
   --rtcp-port=8001       port of the RTCP UDP listener
   --publish-key=""       optional authentication key required to publish
+  --pre-script=""        optional script to run on client connect
+  --post-script=""       optional script to run on client disconnect
 ```
 
 ## Links
diff --git a/client.go b/client.go
index 9727ba6f..f0fb6eae 100644
--- a/client.go
+++ b/client.go
@@ -6,6 +6,7 @@ import (
 	"log"
 	"net"
 	"net/url"
+	"os/exec"
 	"strconv"
 	"strings"
 	"time"
@@ -162,7 +163,18 @@ func (c *client) log(format string, args ...interface{}) {
 }
 
 func (c *client) run() {
+	defer func() {
+                if c.p.postScript != "" {
+                   postScript := exec.Command(c.p.postScript)
+		   err := postScript.Run()
+                   if err != nil {
+			   c.log("ERR: %s", err)
+	           }
+                 }
+        }()
+
 	defer c.log("disconnected")
+
 	defer func() {
 		c.p.mutex.Lock()
 		defer c.p.mutex.Unlock()
@@ -174,6 +186,14 @@ func (c *client) run() {
 
 	c.log("connected")
 
+	if c.p.preScript != "" {
+                   preScript := exec.Command(c.p.preScript)
+		   err := preScript.Run()
+                   if err != nil {
+			   c.log("ERR: %s", err)
+	           }
+        }
+
 	for {
 		req, err := c.conn.ReadRequest()
 		if err != nil {
diff --git a/main.go b/main.go
index 0cdfd13b..6145116e 100644
--- a/main.go
+++ b/main.go
@@ -6,6 +6,7 @@ import (
 	"net"
 	"os"
 	"regexp"
+	"runtime"
 	"strings"
 	"sync"
 	"time"
@@ -53,6 +54,8 @@ type program struct {
 	rtpPort    int
 	rtcpPort   int
 	publishKey string
+	preScript string
+	postScript string
 	mutex      sync.RWMutex
 	rtspl      *serverTcpListener
 	rtpl       *serverUdpListener
@@ -61,7 +64,7 @@ type program struct {
 	publishers map[string]*client
 }
 
-func newProgram(protocolsStr string, rtspPort int, rtpPort int, rtcpPort int, publishKey string) (*program, error) {
+func newProgram(protocolsStr string, rtspPort int, rtpPort int, rtcpPort int, publishKey string, preScript string, postScript string) (*program, error) {
 
 	if rtspPort == 0 {
 		return nil, fmt.Errorf("rtsp port not provided")
@@ -106,6 +109,22 @@ func newProgram(protocolsStr string, rtspPort int, rtpPort int, rtcpPort int, pu
 		}
 	}
 
+	if preScript != "" {
+		if runtime.GOOS != "linux" {
+			return nil, fmt.Errorf("connect script currenty supported only on Linux")
+                } else if !regexp.MustCompile(`(?m)^(.+)\/([^/]+)$`).MatchString(preScript) {
+				return nil, fmt.Errorf("connect script must be a valid path")
+		}
+	}
+
+	if postScript != "" {
+		if runtime.GOOS != "linux" {
+			return nil, fmt.Errorf("disconnect script currently supported only on Linux")
+		} else if !regexp.MustCompile(`(?m)^(.+)\/([^/]+)$`).MatchString(postScript) {
+				return nil, fmt.Errorf("disconnect script must be a valid path")
+		}
+	}
+
 	log.Printf("rtsp-simple-server %s", Version)
 
 	p := &program{
@@ -114,6 +133,8 @@ func newProgram(protocolsStr string, rtspPort int, rtpPort int, rtcpPort int, pu
 		rtpPort:    rtpPort,
 		rtcpPort:   rtcpPort,
 		publishKey: publishKey,
+		preScript:  preScript,
+		postScript: postScript,
 		clients:    make(map[*client]struct{}),
 		publishers: make(map[string]*client),
 	}
@@ -189,6 +210,8 @@ func main() {
 	rtpPort := kingpin.Flag("rtp-port", "port of the RTP UDP listener").Default("8000").Int()
 	rtcpPort := kingpin.Flag("rtcp-port", "port of the RTCP UDP listener").Default("8001").Int()
 	publishKey := kingpin.Flag("publish-key", "optional authentication key required to publish").Default("").String()
+	preScript := kingpin.Flag("pre-script", "optional script to run on client connect").Default("").String()
+	postScript := kingpin.Flag("post-script", "optional script to run on client disconnect").Default("").String()
 
 	kingpin.Parse()
 
@@ -197,7 +220,7 @@ func main() {
 		os.Exit(0)
 	}
 
-	p, err := newProgram(*protocols, *rtspPort, *rtpPort, *rtcpPort, *publishKey)
+	p, err := newProgram(*protocols, *rtspPort, *rtpPort, *rtcpPort, *publishKey, *preScript, *postScript)
 	if err != nil {
 		log.Fatal("ERR: ", err)
 	}