From 7affaf67012eccc6dc9518f2decca790bded59da Mon Sep 17 00:00:00 2001
From: WeebDataHoarder <57538841+WeebDataHoarder@users.noreply.github.com>
Date: Tue, 7 Dec 2021 16:58:18 +0100
Subject: [PATCH] Add support for Postgres output format

---
 SynchRoGazer.go | 75 +++++++++++++++++++++++++++++++++++++++++++------
 go.mod          |  1 +
 go.sum          |  2 ++
 3 files changed, 70 insertions(+), 8 deletions(-)

diff --git a/SynchRoGazer.go b/SynchRoGazer.go
index f03d76a..5aad92b 100644
--- a/SynchRoGazer.go
+++ b/SynchRoGazer.go
@@ -2,14 +2,17 @@ package main
 
 import (
 	"bufio"
+	"database/sql"
 	"encoding/hex"
 	"encoding/json"
 	"flag"
 	"fmt"
+	"github.com/lib/pq"
 	"github.com/minio/md5-simd"
 	"github.com/minio/sha256-simd"
 	"hash"
 	"io"
+	"log"
 	"os"
 	"runtime"
 	"sync/atomic"
@@ -43,9 +46,9 @@ func HashFile(results chan<- HashFileResult, md5hasher *md5simd.Hasher, sha256ha
 	}
 }
 
-func PrintHashFileResult(result *HashFileResult, format string) {
-	switch {
-	case format == "json":
+func PrintHashFileResult(result *HashFileResult, format string, settings PostgresSettings) {
+	switch format {
+	case "json":
 		var jsonData []byte
 		jsonData, err := json.Marshal(*result)
 		if err != nil {
@@ -58,16 +61,48 @@ func PrintHashFileResult(result *HashFileResult, format string) {
 			}
 
 		}
-	case format == "text":
+	case "text":
 		if result.Error != nil {
 			fmt.Fprintln(os.Stderr, result.Path, "Error: ", result.Error)
 		} else {
 			fmt.Println(result.SHA256, result.MD5, result.Path)
 		}
-
+	case "postgres":
+		if result.Error != nil {
+			fmt.Fprintln(os.Stderr, result.Path, "Error: ", result.Error)
+		} else {
+			PostgresHashFileResult(result, settings)
+		}
 	}
 }
 
+func PostgresHashFileResult(result *HashFileResult, settings PostgresSettings) {
+	var err error
+	var rows *sql.Rows
+	switch settings.Mode {
+	case "insert":
+		rows, err = settings.Handle.Query(fmt.Sprintf("INSERT INTO %s (%s, %s, %s) VALUES ($1, $2, $3);", pq.QuoteIdentifier(settings.Table), pq.QuoteIdentifier(settings.PathRow), pq.QuoteIdentifier(settings.MD5Row), pq.QuoteIdentifier(settings.SHA256Row)), result.Path, result.MD5, result.SHA256)
+	case "update":
+		rows, err = settings.Handle.Query(fmt.Sprintf("UPDATE %s SET %s = $2, %s = $3 WHERE %s = $1;", pq.QuoteIdentifier(settings.Table), pq.QuoteIdentifier(settings.MD5Row), pq.QuoteIdentifier(settings.SHA256Row), pq.QuoteIdentifier(settings.PathRow)), result.Path, result.MD5, result.SHA256)
+	}
+	if err != nil {
+		fmt.Fprintln(os.Stderr, result.Path, "SQL Error: ", err)
+	}
+	if rows != nil {
+		rows.Close()
+	}
+}
+
+type PostgresSettings struct {
+	Handle    *sql.DB
+	ConnStr   string
+	Mode      string
+	Table     string
+	PathRow   string
+	MD5Row    string
+	SHA256Row string
+}
+
 func main() {
 	taskLimit := flag.Int("tasklimit", func() int {
 		result := runtime.NumCPU() * 16
@@ -77,10 +112,34 @@ func main() {
 		return result
 	}(), "Maximum number of concurrent hashing tasks. Change to avoid fdlimit issues. Defaults to number of min(128, CPU cores * 16)")
 
-	outputFormat := flag.String("format", "text", "Output formats. Allowed: text, json")
+	outputFormat := flag.String("format", "text", "Output formats. Allowed: text, json, postgres")
+	pgConnStr := flag.String("pg_connstr", "", "Postgres connection string for postgres output format")
+	pgMode := flag.String("pg_mode", "insert", "Postgres output mode, insert or update")
+	pgTable := flag.String("pg_table", "files", "Postgres output table")
+	pgPathRow := flag.String("pg_row_path", "path", "Postgres output row: path")
+	pgMd5Row := flag.String("pg_row_md5", "md5", "Postgres output row: md5")
+	pgSha256Row := flag.String("pg_row_sha256", "sha256", "Postgres output row: sha256")
 
 	flag.Parse()
 
+	pgSettings := PostgresSettings{
+		nil,
+		*pgConnStr,
+		*pgMode,
+		*pgTable,
+		*pgPathRow,
+		*pgMd5Row,
+		*pgSha256Row,
+	}
+
+	if *outputFormat == "postgres" && pgSettings.ConnStr != "" {
+		handle, err := sql.Open("postgres", *pgConnStr)
+		if err != nil {
+			log.Fatal(err)
+		}
+		pgSettings.Handle = handle
+	}
+
 	var taskCount int64
 
 	scanner := bufio.NewScanner(os.Stdin)
@@ -128,13 +187,13 @@ func main() {
 	//Already print before finishing, use atomic ints instead of a WaitGroup
 	for atomic.LoadInt64(&taskCount) > 0 {
 		result := <-resultChannel
-		PrintHashFileResult(&result, *outputFormat)
+		PrintHashFileResult(&result, *outputFormat, pgSettings)
 	}
 
 	close(resultChannel)
 
 	for result := range resultChannel {
-		PrintHashFileResult(&result, *outputFormat)
+		PrintHashFileResult(&result, *outputFormat, pgSettings)
 	}
 
 	close(md5hashers)
diff --git a/go.mod b/go.mod
index b3a8ae3..cb10df7 100644
--- a/go.mod
+++ b/go.mod
@@ -3,6 +3,7 @@ module git.gammaspectra.live/S.O.N.G./SynchRoGazer
 go 1.14
 
 require (
+	github.com/lib/pq v1.10.4
 	github.com/minio/md5-simd v1.1.2
 	github.com/minio/sha256-simd v1.0.0
 )
diff --git a/go.sum b/go.sum
index 2e4b6db..5884eeb 100644
--- a/go.sum
+++ b/go.sum
@@ -1,6 +1,8 @@
 github.com/klauspost/cpuid/v2 v2.0.1/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
 github.com/klauspost/cpuid/v2 v2.0.4 h1:g0I61F2K2DjRHz1cnxlkNSBIaePVoJIjjnHui8QHbiw=
 github.com/klauspost/cpuid/v2 v2.0.4/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
+github.com/lib/pq v1.10.4 h1:SO9z7FRPzA03QhHKJrH5BXA6HU1rS4V2nIVrrNC1iYk=
+github.com/lib/pq v1.10.4/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/minio/md5-simd v1.1.2 h1:Gdi1DZK69+ZVMoNHRXJyNcxrMA4dSxoYHZSQbirFg34=
 github.com/minio/md5-simd v1.1.2/go.mod h1:MzdKDxYpY2BT9XQFocsiZf/NKVtR7nkE4RoEpN+20RM=
 github.com/minio/sha256-simd v1.0.0 h1:v1ta+49hkWZyvaKwrQB8elexRqm6Y0aMLjCNsrYxo6g=