First Commit -- adding files

3 files changed, 236 insertions, 0 deletions

diff --git a/files.go b/files.go
new file mode 100644
index 0000000..def0761
--- /dev/null
+++ b/files.go
@@ -0,0 +1,37 @@
+package main
+
+import (
+	"fmt"
+	"log"
+	"os"
+)
+
+func saveTOfile(filename string, data map[string]datausage) {
+	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
+	if err != nil {
+		log.Fatal(err)
+	}
+	if _, err := fmt.Fprintln(f, "SrcIP Address\t\tRX\t\tTX"); err != nil {
+		log.Fatal(err)
+	}
+	for _, v := range data {
+		if _, err = fmt.Fprintf(f, "%v\t\t%v\t\t%v\n", v.ip, v.rx, v.tx); err != nil {
+			f.Close()
+			log.Fatal(err)
+		}
+	}
+	if err = f.Close(); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func saveTOdatabases(dnm string) {
+	dbx := initDB(dnm)
+	defer dbx.Close()
+	createTable(dbx)
+	storeItem(dbx, nmap)
+	if *svtf != "false" {
+		saveTOfile(*svtf, readItem(dbx))
+	}
+	nmap = make(map[string]datausage)
+}
diff --git a/main.go b/main.go
new file mode 100644
index 0000000..756c448
--- /dev/null
+++ b/main.go
@@ -0,0 +1,78 @@
+package main
+
+import (
+	"flag"
+	"log"
+	"net"
+	"strconv"
+	"time"
+
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/layers"
+	"github.com/google/gopacket/pcapgo"
+)
+
+var finish = make(chan struct{})
+var nmap = make(map[string]datausage)
+var (
+	cidr *net.IPNet
+	svtf *string
+)
+
+func main() {
+	log.SetFlags(0)
+	ipxr := flag.String("net", "192.168.1.0/24", "network to capture on <ipv4/cidr>")
+	infc := flag.String("inf", "lo", "network interface to capture on <interface_name>")
+	dbfi := flag.String("svd", "ipfm.db", "database to save data <database_name>")
+	ftim := flag.String("ttf", "3", "time in second to flush data into the database <integer>")
+	svtf = flag.String("txt", "false", "also save data to file <filename|false>")
+	flag.Parse()
+	if *dbfi == *svtf {
+		log.Fatal("database name and filename can not be the same.")
+	}
+	sectime, err := strconv.Atoi(*ftim)
+	if err != nil {
+		log.Fatal(err)
+	}
+	_, cidr, err = net.ParseCIDR(*ipxr)
+	if err != nil {
+		log.Fatal(err)
+	}
+	handle, err := pcapgo.NewEthernetHandle(*infc)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer handle.Close()
+	log.Printf("starting go-ipfm on %v interface, network %v, database %v", *infc, *ipxr, *dbfi)
+	flush := time.Tick(time.Duration(sectime) * time.Second)
+	packetSource := gopacket.NewPacketSource(handle, layers.LayerTypeEthernet)
+	for packet := range packetSource.Packets() {
+		if ipLayer := packet.Layer(layers.LayerTypeIPv4); ipLayer != nil {
+			ip := ipLayer.(*layers.IPv4)
+			select {
+			case <-flush:
+				saveTOdatabases(*dbfi)
+				accFrom(ip)
+			default:
+				accFrom(ip)
+			}
+		}
+	}
+}
+
+func accFrom(ip *layers.IPv4) {
+	if issrc := cidr.Contains(net.ParseIP(ip.SrcIP.String())); issrc {
+		if val, ok := nmap[ip.SrcIP.String()]; !ok {
+			nmap[ip.SrcIP.String()] = datausage{ip: ip.SrcIP.String(), tx: uint(ip.Length)}
+		} else {
+			nmap[ip.SrcIP.String()] = datausage{ip: ip.SrcIP.String(), tx: uint(ip.Length) + val.tx, rx: val.rx}
+
+		}
+	} else if isdst := cidr.Contains(net.ParseIP(ip.DstIP.String())); isdst {
+		if val, ok := nmap[ip.DstIP.String()]; !ok {
+			nmap[ip.DstIP.String()] = datausage{ip: ip.DstIP.String(), rx: uint(ip.Length)}
+		} else {
+			nmap[ip.DstIP.String()] = datausage{ip: ip.DstIP.String(), rx: uint(ip.Length) + val.rx, tx: val.tx}
+		}
+	}
+}
diff --git a/sqlf.go b/sqlf.go
new file mode 100644
index 0000000..5036d91
--- /dev/null
+++ b/sqlf.go
@@ -0,0 +1,121 @@
+package main
+
+import (
+	"database/sql"
+	"log"
+
+	_ "github.com/mattn/go-sqlite3"
+)
+
+type datausage struct {
+	ip     string
+	tx, rx uint
+}
+
+func initDB(filepath string) *sql.DB {
+	// initial database
+	db, err := sql.Open("sqlite3", filepath)
+	if err != nil {
+		log.Fatal(err)
+	}
+	return db
+}
+
+func createTable(db *sql.DB) {
+	// create table if not exists
+	sqlTable := `
+	CREATE TABLE IF NOT EXISTS items(
+		IP TEXT NOT NULL PRIMARY KEY,
+		RX TEXT,
+		TX TEXT,
+		TIME DATETIME
+	);
+	`
+
+	if _, err := db.Exec(sqlTable); err != nil {
+		log.Fatal(err)
+	}
+
+}
+
+func storeItem(db *sql.DB, items map[string]datausage) {
+	// store items in database
+	sqlAdditem := `
+	INSERT OR REPLACE INTO items(
+		IP,
+		RX,
+		TX,
+		TIME
+	) values(?, ?, ?, CURRENT_TIMESTAMP)
+	`
+
+	sqlSumitem := `
+	INSERT OR REPLACE INTO items(
+		IP,
+		RX,
+		TX,
+		TIME
+	) values(
+		?, 
+		(SELECT RX FROM items WHERE IP=?)+?, 
+		(SELECT TX FROM items WHERE IP=?)+?, 
+		CURRENT_TIMESTAMP)
+	`
+	stmt, err := db.Prepare(sqlAdditem)
+	nstmt, err := db.Prepare(sqlSumitem)
+
+	if err != nil {
+		log.Fatal(err)
+	}
+	for _, item := range items {
+		if ipExists(db, item.ip) {
+			_, err := nstmt.Exec(item.ip, item.ip, item.rx, item.ip, item.tx)
+			if err != nil {
+				log.Fatal(err)
+			}
+			defer nstmt.Close()
+		} else {
+			_, err := stmt.Exec(item.ip, item.rx, item.tx)
+			if err != nil {
+				log.Fatal(err)
+			}
+			defer stmt.Close()
+		}
+	}
+}
+
+func readItem(db *sql.DB) map[string]datausage {
+	// read and return data from database
+	sqlReadall := `
+	SELECT IP, RX, TX FROM items
+	ORDER BY datetime(TIME) DESC
+	`
+	rows, err := db.Query(sqlReadall)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer rows.Close()
+
+	result := make(map[string]datausage)
+	for rows.Next() {
+		item := datausage{}
+		if err := rows.Scan(&item.ip, &item.rx, &item.tx); err != nil {
+			log.Fatal(err)
+		}
+		result[item.ip] = item
+	}
+	return result
+
+}
+
+func ipExists(db *sql.DB, qur string) bool {
+	sqlStmt := `SELECT IP FROM items WHERE IP = ?`
+	err := db.QueryRow(sqlStmt, qur).Scan(&qur)
+	if err != nil {
+		if err != sql.ErrNoRows {
+			log.Fatal(err)
+		}
+		return false
+	}
+	return true
+}