From 42dc6f01d6552015666207d22b423c314313283a Mon Sep 17 00:00:00 2001
From: WeebDataHoarder <57538841+WeebDataHoarder@users.noreply.github.com>
Date: Tue, 18 Jan 2022 20:18:11 +0100
Subject: [PATCH] Added size output postgres entry

---
 .drone.yml      |   2 +-
 .gitignore      |  32 +--------------
 SynchRoGazer.go | 103 +++++++++++++++++++++++++++++++++++++-----------
 go.mod          |   3 +-
 go.sum          |   2 +
 5 files changed, 86 insertions(+), 56 deletions(-)

diff --git a/.drone.yml b/.drone.yml
index fd02f5b..fee4b70 100644
--- a/.drone.yml
+++ b/.drone.yml
@@ -5,7 +5,7 @@ name: default
 
 steps:
   - name: backend
-    image: golang:1.17-bullseye
+    image: golang:1.18-rc-bullseye
     commands:
       - go build -o srg .
 
diff --git a/.gitignore b/.gitignore
index 7e7f1d4..45b9ad9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,5 @@
+/.idea
+
 # ---> Go
 # Binaries for programs and plugins
 *.exe
@@ -15,36 +17,6 @@
 # Dependency directories (remove the comment below to include it)
 # vendor/
 
-# ---> JetBrains
-# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
-# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
-
-# User-specific stuff
-.idea/**/workspace.xml
-.idea/**/tasks.xml
-.idea/**/usage.statistics.xml
-.idea/**/dictionaries
-.idea/**/shelf
-
-# AWS User-specific
-.idea/**/aws.xml
-
-# Generated files
-.idea/**/contentModel.xml
-
-# Sensitive or high-churn files
-.idea/**/dataSources/
-.idea/**/dataSources.ids
-.idea/**/dataSources.local.xml
-.idea/**/sqlDataSources.xml
-.idea/**/dynamic.xml
-.idea/**/uiDesigner.xml
-.idea/**/dbnavigator.xml
-
-# Gradle
-.idea/**/gradle.xml
-.idea/**/libraries
-
 # Gradle and Maven with auto-import
 # When using Gradle or Maven with auto-import, you should exclude module files,
 # since they will be recreated, and may cause churn.  Uncomment if using
diff --git a/SynchRoGazer.go b/SynchRoGazer.go
index 074740f..197f24c 100644
--- a/SynchRoGazer.go
+++ b/SynchRoGazer.go
@@ -23,9 +23,18 @@ type HashFileResult struct {
 	Path   string
 	SHA256 string
 	MD5    string
+	Size   uint64
 }
 
 func HashFile(results chan<- HashFileResult, md5hasher *md5simd.Hasher, sha256hasher *hash.Hash, path string) {
+	fi, err := os.Stat(path)
+	if err != nil {
+		results <- HashFileResult{
+			Error: err,
+			Path:  path,
+		}
+		return
+	}
 	fh, err := os.Open(path)
 	if err != nil {
 		results <- HashFileResult{
@@ -43,6 +52,7 @@ func HashFile(results chan<- HashFileResult, md5hasher *md5simd.Hasher, sha256ha
 		Path:   path,
 		SHA256: hex.EncodeToString((*sha256hasher).Sum(nil)),
 		MD5:    hex.EncodeToString((*md5hasher).Sum(nil)),
+		Size:   uint64(fi.Size()),
 	}
 }
 
@@ -83,16 +93,41 @@ func PostgresHashFileResult(result *HashFileResult, settings PostgresSettings) {
 	hMd5, _ := hex.DecodeString(result.MD5)
 	hSha256, _ := hex.DecodeString(result.SHA256)
 
+	var stmt *sql.Stmt
+
 	switch settings.Mode {
 	case "insert_binary":
-		rows, err = settings.InsertSTMT.Query(result.Path, hMd5, hSha256)
-	case "update_binary":
-		rows, err = settings.UpdateSTMT.Query(result.Path, hMd5, hSha256)
-		break
+		fallthrough
 	case "insert":
-		rows, err = settings.InsertSTMT.Query(result.Path, result.MD5, result.SHA256)
+		stmt = settings.InsertSTMT
+	case "update_binary":
+		fallthrough
 	case "update":
-		rows, err = settings.UpdateSTMT.Query(result.Path, result.MD5, result.SHA256)
+		stmt = settings.UpdateSTMT
+		break
+	}
+
+	if stmt == nil {
+		fmt.Fprintln(os.Stderr, result.Path, "Invalid statement for mode: ", settings.Mode)
+	}
+
+	switch settings.Mode {
+	case "insert_binary":
+		fallthrough
+	case "update_binary":
+		if settings.HasSize {
+			rows, err = stmt.Query(result.Path, hMd5, hSha256, result.Size)
+		} else {
+			rows, err = stmt.Query(result.Path, hMd5, hSha256)
+		}
+	case "insert":
+		fallthrough
+	case "update":
+		if settings.HasSize {
+			rows, err = stmt.Query(result.Path, result.MD5, result.SHA256, result.Size)
+		} else {
+			rows, err = stmt.Query(result.Path, result.MD5, result.SHA256)
+		}
 		break
 	}
 
@@ -107,8 +142,8 @@ func PostgresHashFileResult(result *HashFileResult, settings PostgresSettings) {
 }
 
 type PostgresSettings struct {
-	ConnStr    string
 	Mode       string
+	HasSize    bool
 	InsertSTMT *sql.Stmt
 	UpdateSTMT *sql.Stmt
 }
@@ -129,40 +164,60 @@ func main() {
 	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")
+	pgSizeRow := flag.String("pg_row_size", "", "Postgres output row: size")
 
 	flag.Parse()
 
 	pgSettings := PostgresSettings{
-		*pgConnStr,
-		*pgMode,
-		nil,
-		nil,
+		Mode:    *pgMode,
+		HasSize: len(*pgSizeRow) > 0,
 	}
 
-	if *outputFormat == "postgres" && pgSettings.ConnStr != "" {
+	if *outputFormat == "postgres" && *pgConnStr != "" {
 		handle, err := sql.Open("postgres", *pgConnStr)
 		if err != nil {
 			log.Fatal(err)
 		}
 		defer handle.Close()
 
-		pgSettings.InsertSTMT, err = handle.Prepare(fmt.Sprintf("INSERT INTO %s (%s, %s, %s) VALUES ($1, $2, $3);",
-			pq.QuoteIdentifier(*pgTable),
-			pq.QuoteIdentifier(*pgPathRow),
-			pq.QuoteIdentifier(*pgMd5Row),
-			pq.QuoteIdentifier(*pgSha256Row),
-		))
+		if len(*pgSizeRow) > 0 {
+			pgSettings.InsertSTMT, err = handle.Prepare(fmt.Sprintf("INSERT INTO %s (%s, %s, %s, %s) VALUES ($1, $2, $3, $4);",
+				pq.QuoteIdentifier(*pgTable),
+				pq.QuoteIdentifier(*pgPathRow),
+				pq.QuoteIdentifier(*pgMd5Row),
+				pq.QuoteIdentifier(*pgSha256Row),
+				pq.QuoteIdentifier(*pgSizeRow),
+			))
+		} else {
+			pgSettings.InsertSTMT, err = handle.Prepare(fmt.Sprintf("INSERT INTO %s (%s, %s, %s) VALUES ($1, $2, $3);",
+				pq.QuoteIdentifier(*pgTable),
+				pq.QuoteIdentifier(*pgPathRow),
+				pq.QuoteIdentifier(*pgMd5Row),
+				pq.QuoteIdentifier(*pgSha256Row),
+			))
+		}
 		if err != nil {
 			log.Fatal(err)
 		}
 		defer pgSettings.InsertSTMT.Close()
 
-		pgSettings.UpdateSTMT, err = handle.Prepare(fmt.Sprintf("UPDATE %s SET %s = $2, %s = $3 WHERE %s = $1;",
-			pq.QuoteIdentifier(*pgTable),
-			pq.QuoteIdentifier(*pgPathRow),
-			pq.QuoteIdentifier(*pgMd5Row),
-			pq.QuoteIdentifier(*pgSha256Row),
-		))
+		if len(*pgSizeRow) > 0 {
+			pgSettings.UpdateSTMT, err = handle.Prepare(fmt.Sprintf("UPDATE %s SET %s = $2, %s = $3, %s = $4 WHERE %s = $1;",
+				pq.QuoteIdentifier(*pgTable),
+				pq.QuoteIdentifier(*pgPathRow),
+				pq.QuoteIdentifier(*pgMd5Row),
+				pq.QuoteIdentifier(*pgSha256Row),
+				pq.QuoteIdentifier(*pgSizeRow),
+			))
+		} else {
+			pgSettings.UpdateSTMT, err = handle.Prepare(fmt.Sprintf("UPDATE %s SET %s = $2, %s = $3 WHERE %s = $1;",
+				pq.QuoteIdentifier(*pgTable),
+				pq.QuoteIdentifier(*pgPathRow),
+				pq.QuoteIdentifier(*pgMd5Row),
+				pq.QuoteIdentifier(*pgSha256Row),
+			))
+		}
+
 		if err != nil {
 			log.Fatal(err)
 		}
diff --git a/go.mod b/go.mod
index daa9490..f6a2966 100644
--- a/go.mod
+++ b/go.mod
@@ -1,8 +1,9 @@
 module git.gammaspectra.live/S.O.N.G/SynchRoGazer
 
-go 1.14
+go 1.17
 
 require (
+	github.com/klauspost/cpuid/v2 v2.0.9 // indirect
 	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 5884eeb..49fafe2 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/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4=
+github.com/klauspost/cpuid/v2 v2.0.9/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=