av/cmd/revid-cli/main.go

/*
NAME
  revid-cli - command line interface for Revid.

DESCRIPTION
  See Readme.md

AUTHORS
  Saxon A. Nelson-Milton <saxon@ausocean.org>
  Jack Richardson <jack@ausocean.org>

LICENSE
  revid-cli is Copyright (C) 2017-2018 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It 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 General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.
*/

package main

import (
	"flag"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	"bitbucket.org/ausocean/av/revid"
	"bitbucket.org/ausocean/iot/pi/netsender"
	"bitbucket.org/ausocean/utils/smartlogger"
)

const (
	// progName is the program name for logging purposes.
	progName = "revid-cli"

	// Logging is set to INFO level.
	defaultLogVerbosity = smartlogger.Debug
)

// Other misc consts
const (
	netSendRetryTime   = 5 * time.Second
	defaultRunDuration = 24 * time.Hour
	revidStopTime      = 5 * time.Second
	defaultLogPath     = "/var/log/netsender"
	pkg                = "revid-cli:"
)

// canProfile is set to false with revid-cli is built with "-tags profile".
var canProfile = true

// The logger that will be used throughout
var logger *smartlogger.Logger

func main() {
	useNetsender := flag.Bool("NetSender", false, "Are we checking vars through netsender?")
	runDurationPtr := flag.Duration("runDuration", defaultRunDuration, "How long do you want revid to run for?")

	cfg := handleFlags()

	if !*useNetsender {
		// run revid for the specified duration
		rv, _, err := startRevid(nil, cfg)
		if err != nil {
			cfg.Logger.Log(smartlogger.Fatal, pkg+"failed to start revid", err.Error())
		}
		time.Sleep(*runDurationPtr)
		stopRevid(rv)
		return
	}

	err := run(nil, cfg)
	if err != nil {
		logger.Log(smartlogger.Fatal, pkg+"failed to run revid", "error", err.Error())
		os.Exit(1)
	}
}

// handleFlags parses command line flags and returns a revid configuration
// based on them.
func handleFlags() revid.Config {
	var cfg revid.Config

	var (
		cpuprofile = flag.String("cpuprofile", "", "write cpu profile to `file`")

		inputPtr              = flag.String("Input", "", "The input type: Raspivid, File")
		inputCodecPtr         = flag.String("InputCodec", "", "The codec of the input: H264, Mjpeg")
		output1Ptr            = flag.String("Output1", "", "The first output type: Http, Rtmp, File, Udp, Rtp")
		output2Ptr            = flag.String("Output2", "", "The second output type: Http, Rtmp, File, Udp, Rtp")
		rtmpMethodPtr         = flag.String("RtmpMethod", "", "The method used to send over rtmp: Ffmpeg, Librtmp")
		packetizationPtr      = flag.String("Packetization", "", "The method of data packetisation: Flv, Mpegts, None")
		quantizePtr           = flag.Bool("Quantize", false, "Quantize input (non-variable bitrate)")
		verbosityPtr          = flag.String("Verbosity", "", "Verbosity: Info, Warning, Error, Fatal")
		framesPerClipPtr      = flag.Uint("FramesPerClip", 0, "Number of frames per clip sent")
		rtmpUrlPtr            = flag.String("RtmpUrl", "", "Url of rtmp endpoint")
		bitratePtr            = flag.Uint("Bitrate", 0, "Bitrate of recorded video")
		outputFileNamePtr     = flag.String("OutputFileName", "", "The directory of the output file")
		inputFileNamePtr      = flag.String("InputFileName", "", "The directory of the input file")
		heightPtr             = flag.Uint("Height", 0, "Height in pixels")
		widthPtr              = flag.Uint("Width", 0, "Width in pixels")
		frameRatePtr          = flag.Uint("FrameRate", 0, "Frame rate of captured video")
		httpAddressPtr        = flag.String("HttpAddress", "", "Destination address of http posts")
		quantizationPtr       = flag.Uint("Quantization", 0, "Desired quantization value: 0-40")
		intraRefreshPeriodPtr = flag.Uint("IntraRefreshPeriod", 0, "The IntraRefreshPeriod i.e. how many keyframes we send")
		verticalFlipPtr       = flag.Bool("VerticalFlip", false, "Flip video vertically: Yes, No")
		horizontalFlipPtr     = flag.Bool("HorizontalFlip", false, "Flip video horizontally: Yes, No")
		logPathPtr            = flag.String("LogPath", defaultLogPath, "Path for logging files (default is /var/log/netsender/)")
		rtpAddrPtr            = flag.String("RtpAddr", "", "Rtp destination address: <IP>:<port> (port is generally 6970-6999)")
	)

	flag.Parse()

	logger = smartlogger.New(defaultLogVerbosity, *logPathPtr)

	cfg.Logger = logger

	if *cpuprofile != "" {
		if canProfile {
			f, err := os.Create(*cpuprofile)
			if err != nil {
				logger.Log(smartlogger.Fatal, pkg+"could not create CPU profile", "error", err.Error())
			}
			if err := pprof.StartCPUProfile(f); err != nil {
				logger.Log(smartlogger.Fatal, pkg+"could not start CPU profile", "error", err.Error())
			}
			defer pprof.StopCPUProfile()
		} else {
			logger.Log(smartlogger.Warning, pkg+"ignoring cpuprofile flag - http/pprof built in.")
		}
	}

	switch *inputPtr {
	case "Raspivid":
		cfg.Input = revid.Raspivid
	case "File":
		cfg.Input = revid.File
	case "":
	default:
		logger.Log(smartlogger.Error, pkg+"bad input argument")
	}

	switch *inputCodecPtr {
	case "H264":
		cfg.InputCodec = revid.H264
	case "":
	default:
		logger.Log(smartlogger.Error, pkg+"bad input codec argument")
	}

	switch *output1Ptr {
	case "File":
		cfg.Output1 = revid.File
	case "Http":
		cfg.Output1 = revid.Http
	case "Rtmp":
		cfg.Output1 = revid.Rtmp
	case "FfmpegRtmp":
		cfg.Output1 = revid.FfmpegRtmp
	case "Udp":
		cfg.Output1 = revid.Udp
	case "Rtp":
		cfg.Output1 = revid.Rtp
	case "":
	default:
		logger.Log(smartlogger.Error, pkg+"bad output 1 argument")
	}

	switch *output2Ptr {
	case "File":
		cfg.Output2 = revid.File
	case "Http":
		cfg.Output2 = revid.Http
	case "Rtmp":
		cfg.Output2 = revid.Rtmp
	case "FfmpegRtmp":
		cfg.Output2 = revid.FfmpegRtmp
	case "Udp":
		cfg.Output2 = revid.Udp
	case "Rtp":
		cfg.Output2 = revid.Rtp
	case "":
	default:
		logger.Log(smartlogger.Error, pkg+"bad output 2 argument")
	}

	switch *rtmpMethodPtr {
	case "Ffmpeg":
		cfg.RtmpMethod = revid.Ffmpeg
	case "LibRtmp":
		cfg.RtmpMethod = revid.LibRtmp
	case "":
	default:
		logger.Log(smartlogger.Error, pkg+"bad rtmp method argument")
	}

	switch *packetizationPtr {
	case "None":
		cfg.Packetization = revid.None
	case "Mpegts":
		cfg.Packetization = revid.Mpegts
	case "Flv":
		cfg.Packetization = revid.Flv
	case "":
	default:
		logger.Log(smartlogger.Error, pkg+"bad packetization argument")
	}

	switch *verbosityPtr {
	case "No":
		cfg.LogLevel = smartlogger.Fatal
	case "Debug":
		cfg.LogLevel = smartlogger.Debug
		//logger.SetLevel(smartlogger.Debug)
	case "":
	default:
		logger.Log(smartlogger.Error, pkg+"bad verbosity argument")
	}

	cfg.Quantize = *quantizePtr
	cfg.FlipHorizontal = *horizontalFlipPtr
	cfg.FlipVertical = *verticalFlipPtr
	cfg.FramesPerClip = *framesPerClipPtr
	cfg.RtmpUrl = *rtmpUrlPtr
	cfg.Bitrate = *bitratePtr
	cfg.OutputFileName = *outputFileNamePtr
	cfg.InputFileName = *inputFileNamePtr
	cfg.Height = *heightPtr
	cfg.Width = *widthPtr
	cfg.FrameRate = *frameRatePtr
	cfg.HttpAddress = *httpAddressPtr
	cfg.Quantization = *quantizationPtr
	cfg.IntraRefreshPeriod = *intraRefreshPeriodPtr
	cfg.RtpAddress = *rtpAddrPtr

	return cfg
}

// initialize then run the main NetSender client
func run(rv *revid.Revid, cfg revid.Config) error {
	// initialize NetSender and use NetSender's logger
	//config.Logger = netsender.Logger()
	logger.Log(smartlogger.Info, pkg+"running in NetSender mode")

	var ns netsender.Sender
	err := ns.Init(logger, nil, nil, nil)
	if err != nil {
		return err
	}
	vars, _ := ns.Vars()
	vs := ns.VarSum()
	paused := false
	if vars["mode"] == "Paused" {
		paused = true
	}
	if !paused {
		rv, cfg, err = updateRevid(&ns, rv, cfg, vars, false)
		if err != nil {
			return err
		}
	}

	for {
		if err := send(&ns, rv); err != nil {
			logger.Log(smartlogger.Error, pkg+"polling failed", "error", err.Error())
			time.Sleep(netSendRetryTime)
			continue
		}

		if vs != ns.VarSum() {
			// vars changed
			vars, err := ns.Vars()
			if err != nil {
				logger.Log(smartlogger.Error, pkg+"netSender failed to get vars", "error", err.Error())
				time.Sleep(netSendRetryTime)
				continue
			}
			vs = ns.VarSum()
			if vars["mode"] == "Paused" {
				if !paused {
					logger.Log(smartlogger.Info, pkg+"pausing revid")
					stopRevid(rv)
					paused = true
				}
			} else {
				rv, cfg, err = updateRevid(&ns, rv, cfg, vars, !paused)
				if err != nil {
					return err
				}
				if paused {
					paused = false
				}
			}
		}
		sleepTime, _ := strconv.Atoi(ns.Param("mp"))
		time.Sleep(time.Duration(sleepTime) * time.Second)
	}
}

// send implements our main NetSender client and handles NetReceiver configuration
// (as distinct from httpSender which just sends video data).
func send(ns *netsender.Sender, rv *revid.Revid) error {
	// populate input values, if any
	inputs := netsender.MakePins(ns.Param("ip"), "X")
	if rv != nil {
		for i, pin := range inputs {
			if pin.Name == "X23" {
				inputs[i].Value = rv.Bitrate()
			}
		}
	}

	_, reconfig, err := ns.Send(netsender.RequestPoll, inputs)
	if err != nil {
		return err
	}
	if reconfig {
		return ns.Config()
	}
	return nil
}

// wrappers for stopping and starting revid
func startRevid(ns *netsender.Sender, cfg revid.Config) (*revid.Revid, revid.Config, error) {
	rv, err := revid.New(cfg, ns)
	if err != nil {
		return nil, cfg, err
	}
	rv.Start()
	return rv, cfg, nil
}

func stopRevid(rv *revid.Revid) {
	rv.Stop()

	// FIXME(kortschak): Is this waiting on completion of work?
	// Use a wait group and Wait method if it is.
	time.Sleep(revidStopTime)
}

func updateRevid(ns *netsender.Sender, rv *revid.Revid, cfg revid.Config, vars map[string]string, stop bool) (*revid.Revid, revid.Config, error) {
	if stop {
		stopRevid(rv)
	}

	//look through the vars and update revid where needed
	for key, value := range vars {
		switch key {
		case "Output":
			switch value {
			case "File":
				cfg.Output1 = revid.File
			case "Http":
				cfg.Output1 = revid.Http
			case "Rtmp":
				cfg.Output1 = revid.Rtmp
			case "FfmpegRtmp":
				cfg.Output1 = revid.FfmpegRtmp
			default:
				logger.Log(smartlogger.Warning, pkg+"invalid Output1 param", "value", value)
				continue
			}
		case "FramesPerClip":
			f, err := strconv.ParseUint(value, 10, 0)
			if err != nil {
				logger.Log(smartlogger.Warning, pkg+"invalid framesperclip param", "value", value)
				break
			}
			cfg.FramesPerClip = uint(f)
		case "RtmpUrl":
			cfg.RtmpUrl = value
		case "Bitrate":
			r, err := strconv.ParseUint(value, 10, 0)
			if err != nil {
				logger.Log(smartlogger.Warning, pkg+"invalid framerate param", "value", value)
				break
			}
			cfg.Bitrate = uint(r)
		case "OutputFileName":
			cfg.OutputFileName = value
		case "InputFileName":
			cfg.InputFileName = value
		case "Height":
			h, err := strconv.ParseUint(value, 10, 0)
			if err != nil {
				logger.Log(smartlogger.Warning, pkg+"invalid height param", "value", value)
				break
			}
			cfg.Height = uint(h)
		case "Width":
			w, err := strconv.ParseUint(value, 10, 0)
			if err != nil {
				logger.Log(smartlogger.Warning, pkg+"invalid width param", "value", value)
				break
			}
			cfg.Width = uint(w)
		case "FrameRate":
			r, err := strconv.ParseUint(value, 10, 0)
			if err != nil {
				logger.Log(smartlogger.Warning, pkg+"invalid framerate param", "value", value)
				break
			}
			cfg.FrameRate = uint(r)
		case "HttpAddress":
			cfg.HttpAddress = value
		case "Quantization":
			q, err := strconv.ParseUint(value, 10, 0)
			if err != nil {
				logger.Log(smartlogger.Warning, pkg+"invalid quantization param", "value", value)
				break
			}
			cfg.Quantization = uint(q)
		case "IntraRefreshPeriod":
			p, err := strconv.ParseUint(value, 10, 0)
			if err != nil {
				logger.Log(smartlogger.Warning, pkg+"invalid intrarefreshperiod param", "value", value)
				break
			}
			cfg.IntraRefreshPeriod = uint(p)
		case "HorizontalFlip":
			switch strings.ToLower(value) {
			case "true":
				cfg.FlipHorizontal = true
			case "false":
				cfg.FlipHorizontal = false
			default:
				logger.Log(smartlogger.Warning, pkg+"invalid HorizontalFlip param", "value", value)
			}
		case "VerticalFlip":
			switch strings.ToLower(value) {
			case "true":
				cfg.FlipVertical = true
			case "false":
				cfg.FlipVertical = false
			default:
				logger.Log(smartlogger.Warning, pkg+"invalid VerticalFlip param", "value", value)
			}
		default:
		}
	}

	return startRevid(ns, cfg)
}