KyteAttendance.go

package main

import (
	"encoding/json"
	"flag"
	"errors"
	"fmt"
	"github.com/alexanderkyte/pcap"
	"io/ioutil"
	"log"
	"net/http"
	"html/template"
	"time"
	"encoding/hex"
	"sync"
	"strings"
)

//
//
// Wireless - Pcap Code/Logic
//

const SNAPLEN int32 = 65535
var scanDuration float64
var currentRoom map[string]int

func getNearby(h *pcap.Pcap) map[string]int {
	for {
	ret := make(map[string]int)
	s := time.Now()
	t := time.Now()
	for pkt := h.Next(); pkt != nil && t.Sub(s).Seconds() < scanDuration; pkt = h.Next() {
		pkt.Decode()

		// Only process probe packets. A fraction of the traffic,
		// filters it to just clients.
		// Also skip packets seen as too small to avoid 
		// indexing error. I've checked a few dozen examples by
		// hand and Wireshark doesn't even try to parse them. 
		// It looks like Wireshark discards these probe packets 
		// which are too small, likely interferance or a bug in lib		   // pcap. 
		//
		if pkt.Data[26] != byte(0x40) || len(pkt.Data) < 43 {
		  continue 
		}

		// The radiotap mac address location is located
		// at a different offset than the library wants
		// to think that it is.
		m := pcap.Decodemac(pkt.Data[36:42])

		// for i := 0 ; i < 6 ; i++ {
		// 	fmt.Printf("%x", tmp[i])
		// } 
		// fmt.Printf("\n")
		mac := fmt.Sprintf("%d", m)

		// The signal strength is at offset 8 into
		// the payload. The dB is the negation of
		// 256 minus that value.
		var rssi int = -(256 - int(pkt.Payload[8]))

		// For some reason the pcap library is incorrecly
		// returning the *wireless* mac address. It lies from
		// offsets 36 to 41, inclusive.
		t = time.Now()

		// Omit this check if you want just the mac addresses.
		// Or, add ret[mac] = rssi below to see all results.
		if name, ok := db.Data.Data[mac]; ok {
			// Choose the best signal
			if ret[name] > rssi {
				ret[name] = rssi
			}
		}

	}
	currentRoom = ret
	}
}

//
// Webserver Code/Logic
//
//

func setHeader(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Credentials", "true")
}

func setJsonHeader(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", "GET")
	w.Header().Set("Access-Control-Allow-Credentials", "true")
}

// Format the results into a json api
type Snapshot struct {
	Time   time.Time
	Nearby map[string]int
}

// Display the json api
func showData(w http.ResponseWriter, r *http.Request) {
	s := Snapshot{time.Now(), currentRoom}
	b, err := json.Marshal(s)
	if err != nil {
		log.Println(err)
		return
	}
	fmt.Fprintf(w, string(b))
}

var indexTemplate = template.Must(template.ParseFiles(
  "templates/_base.html",
  "templates/index.html",
))

// Show the list of names and signal strengths
func homePage(w http.ResponseWriter, r *http.Request) {
	setHeader(w)
	if err := indexTemplate.Execute(w, currentRoom); err != nil {
		log.Println(err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

var registrationTemplate = template.Must(template.ParseFiles(
  "templates/_base.html",
  "templates/registration.html",
))

// Show the form for the registration page
func getReg(w http.ResponseWriter, r *http.Request) {
	setHeader(w)

	// Use the template
	if err := registrationTemplate.Execute(w, nil); err != nil {
		log.Println(err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	return
}

// Handle the registration information posted from the registration page
func postReg(w http.ResponseWriter, r *http.Request) {
	setHeader(w)

	// Throw out non-post requests
	if r.Method != "POST" {
		http.Error(w, "Must be a Post Request", http.StatusBadRequest)

	}
	
	// Parse the post params
	err := r.ParseForm()
	if err != nil {
		log.Println(err)
	}
	
	// Get name and mac address
	name := r.FormValue("name")
	mac := r.FormValue("mac")
	mac = strings.Replace(mac, ":", "", -1)

	// Decode the mac address
	macB, err := hex.DecodeString(mac)
	if err != nil {
		http.Error(w, "Invalid MAC Address", http.StatusBadRequest)
	}

	// Using the encoding from libpcap for consistency
	var macC uint64 = pcap.Decodemac(macB) 
	
	err = db.UpdateDB(name, fmt.Sprintf("%d", macC))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
	}
}

//
// Data / Mac-Address-To-User Mapping
//
//

type NameToAddr struct {
	Data map[string]string
}

type DbHandle struct {
	Data     NameToAddr
	Filename string
	Lock     *sync.RWMutex
}

var db DbHandle

// Upon writing to the database, dump the state of the database to
// a json file in the home directory
func (db DbHandle) UpdateDB(name string, mac string) (err error) {
	if len(name) == 0 || len(mac) == 0 {
		return errors.New("Will not process empty fields.")
	}
	db.Lock.Lock()
	db.Data.Data[mac] = name
	db.Lock.Unlock()

	d, err := json.Marshal(db.Data)
	if err != nil {
		return
	}
	err = ioutil.WriteFile(db.Filename, d, 0666)
	if err != nil {
		return
	}
	return nil
}

// Read database from json file
func LoadDb(fileName string) DbHandle {
	db := new(DbHandle)
	db.Filename = fileName
	tmp := new(NameToAddr)
	
	// Read file
	file, err := ioutil.ReadFile(db.Filename)
	if err != nil {
		log.Println("Database file not found. Creating on next save.")
		tmp.Data = make(map[string]string)
	} else if err := json.Unmarshal(file, tmp); err != nil {
		log.Fatal("parse config: ", err)
		panic(err)
	}

	// Initialize the database's lock and information
	db.Data = *tmp
	db.Lock = new(sync.RWMutex)

	return *db
}

//
// Main function loop.
//

// Routes
const (
  homeRoute = "/"
  registrationRoute = "/register"
  processRegistrationRoute = "/processRegistration"
  apiRoute = "/json"
  refreshRoute = "/refresh"
)

func main() {
	// Parse commandline arguments
	var dev = flag.String("dev", "mon0", "The name of the interface to use. Note: this must be the monitor mode interface. Use airmon-ng or something to set it to monitor mode.")
	var dbFile = flag.String("db", "db.json", "The name of the file to use as a database.")
	var port = flag.Int("port", 80, "The port to serve on")
	var scanTime = flag.Int("scan_time", 10, "The number of seconds to scan")
	var pcapFile = flag.String("pcap_file", "", "A pcap file to parse, usually for debugging.")
	flag.Parse()

	scanDuration = float64(*scanTime)
	db = LoadDb(*dbFile)
	currentRoom = make(map[string]int)

	var netDevice *pcap.Pcap
	var err error

	// Either read from file or device
	if *pcapFile == "" {
		// Setting 0 for timeout, the last arg, will
		//  make it never timeout. Also, set
		// promiscuous mode.
		netDevice, err = pcap.OpenLive(*dev, SNAPLEN, true, 0)
	} else {
		netDevice, err = pcap.OpenOffline(*pcapFile)
	}

	if err != nil || netDevice == nil {
		panic(err)
	}

	// Spawn up goroutine for the scanning
	go getNearby(netDevice)


	// Set routes for http server
	http.HandleFunc(homeRoute, homePage)
	http.HandleFunc(registrationRoute, getReg)
	http.HandleFunc(processRegistrationRoute, postReg)
	http.HandleFunc(apiRoute, showData)

	// Output flag
	fmt.Printf("Serving on port %d! Please wait %d seconds to start getting back scans.\n", *port, *scanTime)
	
	// Start serving, make port fit odd go format
	p := fmt.Sprintf(":%d", *port)
	http.ListenAndServe(p, nil)

	return
}