added function and refactored the code

1 files changed, 104 insertions, 0 deletions

diff --git a/sqlQueries.go b/sqlQueries.go
new file mode 100644
index 0000000..4f58242
--- /dev/null
+++ b/sqlQueries.go
@@ -0,0 +1,104 @@
+package main
+
+func fetchRawData(db *sql.DB, numRows int) (rDat []RawData, err error) {
+	prepStmt, err := db.Prepare("SELECT * FROM raw_data LIMIT ? ")
+	if err != nil {
+		log.Println("Failed to prepare statement")
+		return
+	}
+
+	rows, err := prepStmt.Query(numRows)
+	if err != nil {
+		log.Println("Failed to query prepared statement")
+		return
+	}
+
+	loc, err := time.LoadLocation("Local")
+	for rows.Next() {
+		var r RawData
+		var tim []byte
+		err = rows.Scan(&r.ipSrc, &r.ipDst, &tim, &r.port, &r.packetSize)
+		r.time, err = time.ParseInLocation("2006-02-01 15:04:05", string(tim), loc)
+		if err != nil {
+			log.Println("Failed to scan result of query")
+			return
+		}
+		rDat = append(rDat, r)
+	}
+	return
+}
+
+func purgeRawData(tx *sql.Tx, rDat []RawData) (err error) {
+	prepStmt, err := tx.Prepare("DELETE FROM raw_data WHERE ip_src = ? AND ip_dst = ? AND time = ? AND port = ? AND packet_size = ? LIMIT 1")
+	if err != nil {
+		return
+	}
+
+	for _, r := range rDat {
+		_, err = prepStmt.Exec(r.ipSrc, r.ipDst, r.time, r.port, r.packetSize)
+		if err != nil {
+			return
+		}
+	}
+	return
+}
+
+func insertCleanData(tx *sql.Tx, ipbSrc, ipbDst, volume string, time time.Time, port, occurences int) error {
+	prepStmt, err := tx.Prepare("INSERT INTO clean_data (ipb_src, ipb_dst, time, port, volume, occurences) VALUES ( ? , ? , ? , ? , ? , ?) ON DUPLICATE KEY UPDATE occurences = occurences + ?")
+	if err != nil {
+		log.Println("Failed to prepare statement")
+		return err
+	}
+
+	_, err = prepStmt.Exec(ipbSrc, ipbDst, time, port, volume, occurences, occurences)
+	if err != nil {
+		log.Println("Failed to execute statement")
+		return err
+	}
+
+	return nil
+}
+
+func insertASNIP(db *sql.DB, asn int, ipBlock string) error {
+	prepCheck, err := db.Prepare("SELECT asn FROM asnip WHERE ip_block = ?")
+	if err != nil {
+		return err
+	}
+	defer prepCheck.Close()
+
+	rows, err := prepCheck.Exec(ipBlock)
+	if err != nil {
+		return err
+	}
+	if rows != nil {
+		return nil
+	}
+
+	prepIns, err := db.Prepare("INSERT INTO asnip VALUES ( ? , ? )")
+	if err != nil {
+		return err
+	}
+	defer prepIns.Close()
+
+	_, err = prepIns.Exec(asn, ipBlock)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func removeASNIP(db *sql.DB, asn int, ipBlock string) error {
+	prepStmt, err := db.Prepare("DELETE FROM anip WHERE asn = ? AND ip_block = ?")
+	if err != nil {
+		return err
+	}
+	defer prepStmt.Close()
+
+	_, err = prepStmt.Exec(asn, ipBlock)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}