From 1ab69ba1eac11a4ecf359046792f0e7a66945312 Mon Sep 17 00:00:00 2001
From: aceberg <1502200+aceberg@users.noreply.github.com>
Date: Wed, 28 Aug 2024 02:39:41 +0700
Subject: [PATCH] Args for arp-scan

---
 internal/arp/arpscan.go            | 24 ++++++++++++++++++------
 internal/db/choose_db.go           |  3 ++-
 internal/db/postgres.go            | 17 +++++++++++++----
 internal/web/scan-routine.go       |  2 +-
 internal/web/templates/config.html |  2 +-
 5 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/internal/arp/arpscan.go b/internal/arp/arpscan.go
index 518aa01..411dead 100644
--- a/internal/arp/arpscan.go
+++ b/internal/arp/arpscan.go
@@ -6,15 +6,27 @@ import (
 	"strings"
 	"time"
 
+	"github.com/aceberg/WatchYourLAN/internal/check"
 	"github.com/aceberg/WatchYourLAN/internal/models"
 )
 
+var arpArgs string
+
 func scanIface(iface string) string {
-	cmd, err := exec.Command("arp-scan", "-glNx", "-I", iface).Output()
-	if err != nil {
+	var cmd *exec.Cmd
+
+	if arpArgs != "" {
+		cmd = exec.Command("arp-scan", "-glNx", arpArgs, "-I", iface)
+	} else {
+		cmd = exec.Command("arp-scan", "-glNx", "-I", iface)
+	}
+	out, err := cmd.Output()
+	slog.Debug(cmd.String())
+
+	if check.IfError(err) {
 		return string("")
 	}
-	return string(cmd)
+	return string(out)
 }
 
 func parseOutput(text, iface string) []models.Host {
@@ -40,16 +52,16 @@ func parseOutput(text, iface string) []models.Host {
 }
 
 // Scan all interfaces
-func Scan(ifaces string) []models.Host {
+func Scan(ifaces, args string) []models.Host {
 	var text string
 	var foundHosts = []models.Host{}
+	arpArgs = args
 
 	perString := strings.Split(ifaces, " ")
 
 	for _, iface := range perString {
-		text = scanIface(iface)
-
 		slog.Debug("Scanning interface " + iface)
+		text = scanIface(iface)
 		slog.Debug("Found IPs: \n" + text)
 
 		foundHosts = append(foundHosts, parseOutput(text, iface)...)
diff --git a/internal/db/choose_db.go b/internal/db/choose_db.go
index dbb2a95..4ca05a9 100644
--- a/internal/db/choose_db.go
+++ b/internal/db/choose_db.go
@@ -44,9 +44,10 @@ func SetCurrent(config models.Conf) {
 	currentDB.SQLitePath = config.DBPath
 	currentDB.PGConnect = config.PGConnect
 
-	if currentDB.Use == "postgres" {
+	if currentDB.Use == "postgres" && currentDB.PGConnect != "" {
 		currentDB.PrimaryKey = "BIGSERIAL PRIMARY KEY"
 	} else {
+		currentDB.Use = "sqlite"
 		currentDB.PrimaryKey = "INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT UNIQUE"
 	}
 
diff --git a/internal/db/postgres.go b/internal/db/postgres.go
index 233a7ab..a2763ef 100644
--- a/internal/db/postgres.go
+++ b/internal/db/postgres.go
@@ -1,7 +1,7 @@
 package db
 
 import (
-	// "log/slog"
+	"log/slog"
 
 	"github.com/jmoiron/sqlx"
 
@@ -17,7 +17,11 @@ func dbExecPG(sqlStatement string) {
 	// slog.Debug("PG Exec " + sqlStatement)
 
 	db, err := sqlx.Connect("postgres", currentDB.PGConnect)
-	check.IfError(err)
+	if check.IfError(err) {
+		slog.Warn("PostgreSQL connection error. Falling back to SQLite.")
+		currentDB.Use = "sqlite"
+		Create()
+	}
 	defer db.Close()
 
 	_, err = db.Exec(sqlStatement)
@@ -30,10 +34,15 @@ func selectPG(table string) (dbHosts []models.Host) {
 
 	// slog.Debug("PG Select " + sqlStatement)
 
-	db, _ := sqlx.Connect("postgres", currentDB.PGConnect)
+	db, err := sqlx.Connect("postgres", currentDB.PGConnect)
+	if check.IfError(err) {
+		slog.Warn("PostgreSQL connection error. Falling back to SQLite.")
+		currentDB.Use = "sqlite"
+		Create()
+	}
 	defer db.Close()
 
-	err := db.Select(&dbHosts, sqlStatement)
+	err = db.Select(&dbHosts, sqlStatement)
 	check.IfError(err)
 
 	return dbHosts
diff --git a/internal/web/scan-routine.go b/internal/web/scan-routine.go
index 35b15ec..ca82e94 100644
--- a/internal/web/scan-routine.go
+++ b/internal/web/scan-routine.go
@@ -26,7 +26,7 @@ func startScan(quit chan bool) {
 
 			if nowDate.After(plusDate) {
 
-				foundHosts = arp.Scan(appConfig.Ifaces)
+				foundHosts = arp.Scan(appConfig.Ifaces, appConfig.ArpArgs)
 				compareHosts(foundHosts)
 				allHosts = db.Select("now")
 
diff --git a/internal/web/templates/config.html b/internal/web/templates/config.html
index 5d4f1c0..003468d 100644
--- a/internal/web/templates/config.html
+++ b/internal/web/templates/config.html
@@ -78,7 +78,7 @@
                   <td><input name="timeout" type="number" class="form-control" value="{{ .Config.Timeout }}"></td>
                 </tr>
                 <tr>
-                  <td>Args for arpscan</td>
+                  <td>Args for arp-scan</td>
                   <td><input name="arpargs" type="text" class="form-control" value="{{ .Config.ArpArgs }}"></td>
                 </tr>
                 <tr>