nmdc-log-service/main.go

package main

import (
	"flag"
	"fmt"
	"libnmdc"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

var BaseDir string = "."
var PMResponse string = ""
var LogConnectionState bool
var DebugMode bool
var VerifyTLS bool
var CharacterMatcher *regexp.Regexp

func init() {
	CharacterMatcher = regexp.MustCompile("[^a-zA-Z0-9]")
}

func GetDirectoryNameForHub(hub string) string {
	return filepath.Join(BaseDir, CharacterMatcher.ReplaceAllString(hub, "_"))
}

func LogMessage(hub, message string) {

	instant := time.Now() // using system timezone

	// Log file path
	monthStamp := instant.Format("2006-01")
	logFilePath := filepath.Join(GetDirectoryNameForHub(hub), monthStamp+".log")

	// Produce a timestamp
	timeStamp := instant.Format("2006-01-02 15:04:05")

	fh, err := os.OpenFile(logFilePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't open file '%s': %s\n", logFilePath, err.Error())
		return
	}
	defer fh.Close()

	_, err = fh.WriteString("[" + timeStamp + "] " + message + "\n")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error writing to file '%s': %s\n", logFilePath, err.Error())
		return
	}
}

type HubWorker struct {
	LastConnectionStateMessage string
}

func (this *HubWorker) MainLoop(addr, nick, password string) {

	opts := libnmdc.HubConnectionOptions{
		Address:       libnmdc.HubAddress(addr),
		SkipVerifyTLS: !VerifyTLS,
		Self:          *libnmdc.NewUserInfo(nick),
		NickPassword:  password,
	}
	hub := opts.Connect()

	for {
		event := <-hub.OnEvent

		if DebugMode {
			fmt.Printf("DEBUG: %s %v\n", addr, event)
		}

		switch event.EventType {
		case libnmdc.EVENT_CONNECTION_STATE_CHANGED:
			if LogConnectionState {
				str := "* " + event.StateChange.Format()
				if len(event.Message) > 0 {
					str += " (" + event.Message + ")"
				}

				// Prevent logging the same message repeatedly
				if str != this.LastConnectionStateMessage {
					LogMessage(addr, str)
				}
				this.LastConnectionStateMessage = str
			}

		case libnmdc.EVENT_PUBLIC:
			LogMessage(addr, "<"+event.Nick+"> "+event.Message)

		case libnmdc.EVENT_PRIVATE:
			fmt.Printf("Got PM %v\n", event)
			hub.SayPrivate(event.Nick, PMResponse)

		case libnmdc.EVENT_SYSTEM_MESSAGE_FROM_CONN, libnmdc.EVENT_SYSTEM_MESSAGE_FROM_HUB:
			if strings.HasPrefix(event.Message, "* ") {
				LogMessage(addr, event.Message)
			} else {
				LogMessage(addr, "* "+event.Message)
			}

		}
	}
}

func main() {

	// Parse arguments
	hubs := flag.String("Server", "", "Addresses to connect to (comma-separated)")
	nick := flag.String("Nick", "nmdc-log-service", "Nick")
	password := flag.String("Password", "", "Registered nick password")
	flag.StringVar(&BaseDir, "Dir", ".", "Output directory")
	flag.BoolVar(&LogConnectionState, "LogConnectionState", true, "Include connection state changes in log")
	flag.StringVar(&PMResponse, "PMResponse", "This is an automated service. For enquiries, please contact an administrator.", "Message to respond with on PM")
	flag.BoolVar(&DebugMode, "Debug", false, "Print additional information on stdout")
	flag.BoolVar(&VerifyTLS, "VerifyTLS", true, "Verify TLS certificates")
	flag.Parse()

	// Assert dir exists
	dinfo, err := os.Stat(BaseDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "FATAL: %s\n", err.Error())
		os.Exit(1)
	}
	if !dinfo.IsDir() {
		fmt.Fprintf(os.Stderr, "FATAL: Path '%s' is not a directory\n", BaseDir)
	}

	if !VerifyTLS {
		fmt.Fprintf(os.Stderr, "WARNING: TLS certificates will not be verified!\n")
	}

	// Launch loggers
	all_hubs := strings.Split(*hubs, ",")
	launch_ct := 0
	for _, hubaddr := range all_hubs {
		if len(hubaddr) == 0 {
			continue
		}

		// Assert target directory exists
		os.MkdirAll(GetDirectoryNameForHub(hubaddr), 0755)

		// Launch logger
		hw := HubWorker{}
		go hw.MainLoop(hubaddr, *nick, *password)

		launch_ct++
	}

	if launch_ct == 0 {
		fmt.Fprintln(os.Stderr, "FATAL: No hubs specified")
		os.Exit(1)
	}

	// Wait forever
	var forever chan bool = nil
	select {
	case <-forever:
	}
}