From bf9ad497ce38be60b8387ea9dab932c40b2ca372 Mon Sep 17 00:00:00 2001
From: bn0ir <gblacknoir@gmail.com>
Date: Wed, 17 Apr 2019 18:31:36 +0500
Subject: [PATCH] Add postgresql audit backend

---
 audit.go       | 75 ++++++++++++++++++++++++++++++++++++++++++++++++++
 etc/godns.conf | 13 ++++++++-
 settings.go    | 29 +++++++++++++------
 3 files changed, 108 insertions(+), 9 deletions(-)

diff --git a/audit.go b/audit.go
index 09f10c7..3a5858f 100644
--- a/audit.go
+++ b/audit.go
@@ -1,11 +1,13 @@
 package main
 
 import (
+	"database/sql"
 	"encoding/json"
 	"fmt"
 	"time"
 
 	"github.com/hoisie/redis"
+	_ "github.com/lib/pq"
 )
 
 const AUDIT_LOG_OUTPUT_BUFFER = 1024
@@ -75,3 +77,76 @@ func (rl *RedisAuditLogger) Run() {
 func (rl *RedisAuditLogger) Write(mesg *AuditMesg) {
 	rl.mesgs <- mesg
 }
+
+type PostgresqlAuditLogger struct {
+	backend *sql.DB
+	mesgs   chan *AuditMesg
+	expire  int64
+}
+
+func NewPostgresqlAuditLogger(ps PostgresqlSettings, expire int64) AuditLogger {
+	connStr := fmt.Sprintf(`
+                host=%s port=%d
+                user=%s password=%s
+                dbname=%s sslmode=%s
+                sslcert=%s sslkey=%s
+                sslrootcert=%s
+                `,
+		ps.Host, ps.Port,
+		ps.User, ps.Password,
+		ps.DB, ps.Sslmode,
+		ps.Sslcert, ps.Sslkey,
+		ps.Sslrootcert,
+	)
+	pc, err := sql.Open("postgres", connStr)
+	if err != nil {
+		logger.Error("Can't connect to audit log postgresql: %v", err)
+	}
+	_, err = pc.Query(`
+                CREATE TABLE IF NOT EXISTS audit (
+                        id BIGSERIAL NOT NULL,
+                        remoteaddr TEXT,
+                        domain TEXT,
+                        qtype TEXT,
+                        timestamp TIMESTAMP
+                )
+        `)
+	auditLogger := &PostgresqlAuditLogger{
+		backend: pc,
+		mesgs:   make(chan *AuditMesg, AUDIT_LOG_OUTPUT_BUFFER),
+		expire:  expire,
+	}
+	go auditLogger.Run()
+	go auditLogger.Expire()
+	return auditLogger
+}
+
+func (pl *PostgresqlAuditLogger) Run() {
+	for {
+		select {
+		case mesg := <-pl.mesgs:
+			_, err := pl.backend.Query(`INSERT INTO audit (remoteaddr, domain, qtype, timestamp) VALUES ($1, $2, $3, $4)`,
+				mesg.RemoteAddr, mesg.Domain, mesg.QType, mesg.Timestamp,
+			)
+			if err != nil {
+				logger.Error("Can't write to postgresql audit log: %v", err)
+				continue
+			}
+		}
+	}
+}
+
+func (pl *PostgresqlAuditLogger) Write(mesg *AuditMesg) {
+	pl.mesgs <- mesg
+}
+
+func (pl *PostgresqlAuditLogger) Expire() {
+	for {
+		expireTime := time.Now().Add(time.Duration(-pl.expire) * time.Second)
+		_, err := pl.backend.Query(`DELETE FROM audit WHERE timestamp < $1`, expireTime)
+		if err != nil {
+			logger.Error("Can't expire postgresql audit log: %v", err)
+		}
+		time.Sleep(time.Duration(pl.expire) * time.Second / 2)
+	}
+}
diff --git a/etc/godns.conf b/etc/godns.conf
index 059dea8..76b44e7 100644
--- a/etc/godns.conf
+++ b/etc/godns.conf
@@ -30,6 +30,17 @@ port = 6379
 db = 0
 password =""
 
+[postgresql]
+host = "127.0.0.1"
+port = 5432
+user = "godns"
+password = "godns"
+db = "godns"
+sslmode = "require"
+sslcert = ""
+sslkey = ""
+sslrootcert = ""
+
 [memcache]
 servers = ["127.0.0.1:11211"]
 
@@ -40,7 +51,7 @@ level = "INFO"  #DEBUG | INFO |NOTICE | WARN | ERROR
 
 
 # [audit]
-# backend option [redis|file]
+# backend option [redis|postgresql|file]
 # file backend not implemented yet
 # backend = "redis"
 # expire = 864000 # 10 days
diff --git a/settings.go b/settings.go
index acc9257..9e30e2b 100644
--- a/settings.go
+++ b/settings.go
@@ -24,14 +24,15 @@ var LogLevelMap = map[string]int{
 type Settings struct {
 	Version      string
 	Debug        bool
-	Server       DNSServerSettings `toml:"server"`
-	ResolvConfig ResolvSettings    `toml:"resolv"`
-	Redis        RedisSettings     `toml:"redis"`
-	Memcache     MemcacheSettings  `toml:"memcache"`
-	Log          LogSettings       `toml:"log"`
-	Cache        CacheSettings     `toml:"cache"`
-	Hosts        HostsSettings     `toml:"hosts"`
-	Audit        AuditSettings     `toml:"audit"`
+	Server       DNSServerSettings  `toml:"server"`
+	ResolvConfig ResolvSettings     `toml:"resolv"`
+	Redis        RedisSettings      `toml:"redis"`
+	Memcache     MemcacheSettings   `toml:"memcache"`
+	Postgresql   PostgresqlSettings `toml:"postgresql"`
+	Log          LogSettings        `toml:"log"`
+	Cache        CacheSettings      `toml:"cache"`
+	Hosts        HostsSettings      `toml:"hosts"`
+	Audit        AuditSettings      `toml:"audit"`
 }
 
 type ResolvSettings struct {
@@ -54,6 +55,18 @@ type RedisSettings struct {
 	Password string
 }
 
+type PostgresqlSettings struct {
+	Host        string
+	Port        int
+	DB          string
+	User        string
+	Password    string
+	Sslmode     string
+	Sslcert     string
+	Sslkey      string
+	Sslrootcert string
+}
+
 type MemcacheSettings struct {
 	Servers []string
 }