nim-cli-utils/cliutils.nim

import docopt, json, osproc, posix, nre, streams, strtabs, tables, terminal, unicode
import os except sleep
import strutils except toUpper, toLower


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

  TermColor = ForegroundColor or BackgroundColor

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):
    let node = cfg.json[jsonKey]
    case node.kind
    of JString: return node.getStr
    of JInt: return $node.getInt
    of JFloat: return $node.getFloat
    of JBool: return $node.getBool
    of JNull: return ""
    of JObject: return $node
    of JArray: return $node
  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 = "", cmd: string = ""): void =
  if h != nil: h(outMsg, errMsg, cmd)

proc makeProcMsgHandler*(outSink, errSink: File, prefixCmd: bool = true): HandleProcMsgCB =
  result = proc(outMsg, errMsg: TaintedString, cmd: string): void {.closure.} =
    let prefix = if cmd.len == 0 or not prefixCmd: "" else: cmd & ": "
    if outMsg.len > 0: outSink.writeLine(prefix & outMsg)
    if errMsg.len > 0: errSink.writeLine(prefix & errMsg)

proc makeProcMsgHandler*(outSink, errSink: Stream, prefixCmd: bool = true): HandleProcMsgCB =
  result = proc(outMsg, errMsg: TaintedString, cmd: string): void {.closure.} =
    let prefix = if cmd.len == 0 or not prefixCmd: "" else: cmd & ": "
    if outMsg.len > 0: outSink.writeLine(prefix & outMsg)
    if errMsg.len > 0: 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)

proc waitFor*(p: Process, msgCB: HandleProcMsgCB, procCmd: string = ""): int =

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

  var line = newStringOfCap(120).TaintedString
  while true:
    if pout.readLine(line):
      msgCB.sendMsg(line, "", procCmd)
    elif perr.readLine(line):
      msgCB.sendMsg("", line, procCmd)
    else:
      result = peekExitCode(p)
      if result != -1: break
  close(p)

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

  var p = startProcess(command, workingDir, args, env, options)
  result = waitFor(p, msgCB, command)

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

  result = (TaintedString"", TaintedString"", -1)
  var outSeq, errSeq: seq[TaintedString]
  outSeq = @[]; errSeq = @[]
  var outputCollector = combineProcMsgHandlers(msgCB,
    proc(outMsg, errMsg: TaintedString, cmd: string): void {.closure.} =
      if outMsg.len > 0: outSeq.add(outMsg)
      if errMsg.len > 0: errSeq.add(errMsg))

  result[2] = exec(command, workingDir, args, env, options, outputCollector)
  result[0] = outSeq.join("\n")
  result[1] = errSeq.join("\n")

## 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): Pid =

  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 si.len > 0:
      fi = open(si, fmRead)
      discard dup2(getFileHandle(fi), getFileHandle(stdin))

    if so.len > 0:
      fo = open(so, fmAppend)
      discard dup2(getFileHandle(fo), getFileHandle(stdout))

    if se.len > 0:
      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

const termReset* = "\e[0;m"

proc termColor*(color: TermColor, bright, bold = false): string =
  var colorVal = ord(color)
  if bright: inc(colorVal, 60)
  return "\e[" & $colorVal & (if bold: ";1" else: "") & "m"

proc withColor*(str: string, color: TermColor, bright, bold = false): string =
  return termColor(color, bright, bold) & str

let STRIP_ANSI_REGEX = re"\x1B\[([0-9]{1,2}(;[0-9]{1,2})?)?[m|K]"

proc stripAnsi*(str: string): string = return str.replace(STRIP_ANSI_REGEX, "")

proc queryParamsToCliArgs*(queryParams: Table[string, string]): seq[string] =
  result = @[]

  for k,v in queryParams:
    # support ?arg1=val1&arg2=val2 -> cmd val1 val2
    if k.startsWith("arg"): result.add(v)

    else :
      result[1].add("--" & k)
      if v != "true": result[1].add(v) # support things like ?verbose=true -> cmd --verbose

proc queryParamsToCliArgs*(queryParams: StringTableRef): seq[string] =
  result = @[]

  for k,v in queryParams:
    # support ?arg1=val1&arg2=val2 -> cmd val1 val2
    if k.startsWith("arg"): result.add(v)

    else :
      result[1].add("--" & k)
      if v != "true": result[1].add(v) # support things like ?verbose=true -> cmd --verbose