nim-cli-utils/cliutils.nim

import docopt, json, osproc, posix, nre, streams, strutils, strtabs
import os except sleep

type
  CombinedConfig* = object
    docopt*: Table[string, Value]
    json*: JsonNode

proc getVal*(cfg: CombinedConfig, key, default: string): string =
  let argKey = "--" & key
  let envKey = key.replace('-', '_').toUpper
  let jsonKey = key.replace(re"(-\w)", proc (m: RegexMatch): string = ($m)[1..1].toUpper)

  if cfg.docopt[argKey]: return $cfg.docopt[argKey]
  elif existsEnv(envKey): return getEnv(envKey)
  elif cfg.json.hasKey(jsonKey): return cfg.json[jsonKey].getStr
  else: return default


proc loadEnv*(): StringTableRef =
  result = newStringTable()

  for k, v in envPairs():
    result[k] = v

## Process execution
type HandleProcMsgCB* = proc (outMsg: TaintedString,
                              errMsg: TaintedString, cmd: string): void

proc sendMsg*(h: HandleProcMsgCB, outMsg: TaintedString,
              errMsg: TaintedString = nil, cmd: string): void =
  if h != nil: h(outMsg, errMsg, cmd)

proc waitForWithOutput*(p: Process, msgCB: HandleProcMsgCB,
                        procCmd: string = ""):
      tuple[output: TaintedString, error: TaintedString, exitCode: int] =

  var pout = outputStream(p)
  var perr = errorStream(p)

  result = (TaintedString"", TaintedString"", -1)
  var line = newStringOfCap(120).TaintedString
  while true:
    if pout.readLine(line):
      msgCB.sendMsg(line, nil, procCmd)
      result[0].string.add(line.string)
      result[0].string.add("\n")
    elif perr.readLine(line):
      msgCB.sendMsg(nil, line, procCmd)
      result[1].string.add(line.string)
      result[1].string.add("\n")
    else:
      result[2] = peekExitCode(p)
      if result[2] != -1: break
  close(p)

proc exec*(command: string, workingDir: string = "",
           args: openArray[string] = [], env: StringTableRef = nil,
           options: set[ProcessOption] = {poUsePath},
           msgCB: HandleProcMsgCB = nil):
    tuple[output: TaintedString, error: TaintedString, exitCode: int]
    {.tags: [ExecIOEffect, ReadIOEffect], gcsafe.} =

  var p = startProcess(command, workingDir, args, env, options)
  result = waitForWithOutput(p, msgCb, command)

proc makeProcMsgHandler*(outSink, errSink: File): HandleProcMsgCB =
  result = proc(outMsg, errMsg: TaintedString, cmd: string): void {.closure.} =
    let prefix = if cmd != nil: cmd & ": " else: ""
    if outMsg != nil: outSink.writeLine(prefix & outMsg)
    if errMsg != nil: errSink.writeLine(prefix & errMsg)

proc makeProcMsgHandler*(outSink, errSink: Stream): HandleProcMsgCB =
  result = proc(outMsg, errMsg: TaintedString, cmd: string): void {.closure.} =
    let prefix = if cmd != nil: cmd & ": " else: ""
    if outMsg != nil: outSink.writeLine(prefix & outMsg)
    if errMsg != nil: errSink.writeLine(prefix & errMsg)

proc combineProcMsgHandlers*(a, b: HandleProcMsgCB): HandleProcMsgCB =
  if a == nil: result = b
  elif b == nil: result = a
  else:
    result = proc(cmd: string, outMsg, errMsg: TaintedString): void =
      a(cmd, outMsg, errMsg)
      b(cmd, outMsg, errMsg)


## Daemonize

var
  pidFileInner: string
  fi, fo, fe: File

proc onStop(sig: cint) {.noconv.} =
  close(fi)
  close(fo)
  close(fe)
  removeFile(pidFileInner)

  quit(QuitSuccess)

proc daemonize*(pidfile, si, so, se: string, daemonMain: proc(): void): int =

  if fileExists(pidfile):
    raise newException(IOError, "pidfile " & pidfile & " already exists, daemon already running?")

  let pid1 = fork()

  # Are we the child process?
  if pid1 == 0:

    # Yes, so let's get ready to execute the main code given to us.
    discard chdir("/")
    discard setsid()
    discard umask(0)

    # Fork again... but I'm not sure why.
    let pid2 = fork()

    # We don't need the intermediate process, so if we are not the child this
    # time, lets just quit
    if pid2 > 0: quit(QuitSuccess)

    # If we are the grandchild let's set up our environment.
    flushFile(stdout)
    flushFile(stderr)

    if not si.isNil and si != "":
      fi = open(si, fmRead)
      discard dup2(getFileHandle(fi), getFileHandle(stdin))

    if not so.isNil and so != "":
      fo = open(so, fmAppend)
      discard dup2(getFileHandle(fo), getFileHandle(stdout))

    if not se.isNil and so != "":
      fe = open(se, fmAppend)
      discard dup2(getFileHandle(fe), getFileHandle(stderr))

    pidFileInner = pidfile

    # Add hooks to cleanup after ourselves when we're asked to die.
    signal(SIGINT, onStop)
    signal(SIGTERM, onStop)

    # Find out what our actual PID is and save it
    let childPid = getpid()
    writeFile(pidfile, $childPid)

    # Finally, execute our main
    daemonMain()

  return pid1