From a399a1cc8df6edab1a2bd750f1d24d6ebcc262de Mon Sep 17 00:00:00 2001
From: IHaxYou <andrewcscott007@gmail.com>
Date: Fri, 19 Jan 2018 23:34:37 -0700
Subject: [PATCH] extended checks to src/dst for blacklist

---
 cli/sql.go  |  6 +++---
 firewall.go | 31 +++++++++++++++++++++----------
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/cli/sql.go b/cli/sql.go
index 32f7ac1..68dfde8 100644
--- a/cli/sql.go
+++ b/cli/sql.go
@@ -26,7 +26,7 @@ func sqlCheckExists(table string) bool {
 func init() {
 	// open database instance
 	var err error
-	db, err = sql.Open("sqlite3", "db/database.db")
+	db, err = sql.Open("sqlite3", "goaway2/db/database.db")
 	if err != nil {
 		log.Fatalf("Unable to launch SQLITE3: %s\n", err.Error())
 	}
@@ -34,9 +34,9 @@ func init() {
 	db.SetMaxOpenConns(1)
 	db.Exec("PRAGMA journal_mode=WAL;")
 	// get reusable sql functions
-	dot, err := dotsql.LoadFromFile("tables.sql")
+	dot, err := dotsql.LoadFromFile("goaway2/tables.sql")
 	if err != nil {
-		log.Fatalln(err.Error())
+		log.Fatalf("Unable to load SQL: %s\n", err.Error())
 	}
 	// check if required tables exist
 	if !sqlCheckExists("rules") {
diff --git a/firewall.go b/firewall.go
index 3795fc8..3671e33 100644
--- a/firewall.go
+++ b/firewall.go
@@ -36,11 +36,15 @@ func NewFirewall() *Firewall {
 /***Methods***/
 
 //(*Firewall).HandlePackets : packet hander used to block/allow packets based on rules
-func (fw *Firewall) HandlePackets(kv *RBKV, pkt *PacketData) netfilter.Verdict {
+func (fw *Firewall) HandlePackets(l *log.Logger, kv *RBKV, pkt *PacketData) netfilter.Verdict {
 	switch {
 	// if src-ip is in blacklist cache
 	case fw.blacklist.Exists(kv, pkt.SrcIP):
-		log.Printf("Fast Block: %s\n", pkt.SrcIP)
+		l.Printf("Fast Block SRC: %s\n", pkt.SrcIP)
+		return netfilter.NF_DROP
+	// if dst-ip is in blacklist cache
+	case fw.blacklist.Exists(kv, pkt.DstIP):
+		l.Printf("Fast Block DST: %s\n", pkt.SrcIP)
 		return netfilter.NF_DROP
 	// if src-ip is in whitelist cache
 	case fw.whitelist.Exists(kv, pkt.SrcIP):
@@ -50,16 +54,23 @@ func (fw *Firewall) HandlePackets(kv *RBKV, pkt *PacketData) netfilter.Verdict {
 		return fw.checkRules(pkt)
 	// if src-ip is not in a cache
 	default:
-		var blocked int
-		db.QueryRow("SELECT EXISTS(SELECT 1 FROM blacklist WHERE LogicalDelete=0 AND IPAddress=?)", pkt.SrcIP).Scan(&blocked)
-		// if they are blocked, add to cache and drop packet
-		if blocked == 1 {
+		var blocked string
+		db.QueryRow("SELECT IPAddress FROM blacklist WHERE LogicalDelete=0 AND (IPAddress=? OR IPAddress=?)", pkt.SrcIP, pkt.DstIP).Scan(&blocked)
+		switch blocked {
+		case pkt.SrcIP:
+			// if source ip is blacklisted
 			fw.blacklist.Set(kv, pkt.SrcIP, "")
 			return netfilter.NF_DROP
+		case pkt.DstIP:
+			// if destination ip is blacklisted
+			fw.blacklist.Set(kv, pkt.DstIP, "")
+			return netfilter.NF_DROP
+		default:
+			// else put them in the neutral cache and evaluate the rules
+			fw.neutlist.Set(kv, pkt.SrcIP, "")
+			fw.neutlist.Set(kv, pkt.DstIP, "")
+			return fw.checkRules(pkt)
 		}
-		// else put them in the neutral cache and evaluate the rules
-		fw.neutlist.Set(kv, pkt.SrcIP, "")
-		return fw.checkRules(pkt)
 	}
 }
 
@@ -99,7 +110,7 @@ func (fw *Firewall) checkRules(pkt *PacketData) netfilter.Verdict {
 				} else {
 					continue
 				}
-			// if outbounds deafault is to deny
+			// if outbounds default is to deny
 			default:
 				// if rule matches: accept
 				if rule.Validate(pkt) {