Add support for Postgres output format
This commit is contained in:
parent
b847a14bce
commit
7affaf6701
|
|
@ -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)
|
||||
|
|
|
|||
1
go.mod
1
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
|
||||
)
|
||||
|
|
|
|||
2
go.sum
2
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=
|
||||
|
|
|
|||
Loading…
Reference in a new issue