wifi_ap.go 2.83 KB
Newer Older
Sophie Brun's avatar
Sophie Brun committed
1 2 3 4 5 6
package network

import (
	"encoding/json"
	"sync"
	"time"
Sophie Brun's avatar
Sophie Brun committed
7 8

	"github.com/evilsocket/islazy/data"
Sophie Brun's avatar
Sophie Brun committed
9 10 11 12 13 14
)

type AccessPoint struct {
	*Station
	sync.Mutex

Sophie Brun's avatar
Sophie Brun committed
15
	aliases         *data.UnsortedKV
Sophie Brun's avatar
Sophie Brun committed
16 17
	clients         map[string]*Station
	withKeyMaterial bool
Sophie Brun's avatar
Sophie Brun committed
18 19 20 21
}

type apJSON struct {
	*Station
Sophie Brun's avatar
Sophie Brun committed
22 23
	Clients   []*Station `json:"clients"`
	Handshake bool       `json:"handshake"`
Sophie Brun's avatar
Sophie Brun committed
24 25
}

Sophie Brun's avatar
Sophie Brun committed
26
func NewAccessPoint(essid, bssid string, frequency int, rssi int8, aliases *data.UnsortedKV) *AccessPoint {
Sophie Brun's avatar
Sophie Brun committed
27 28
	return &AccessPoint{
		Station: NewStation(essid, bssid, frequency, rssi),
Sophie Brun's avatar
Sophie Brun committed
29
		aliases: aliases,
Sophie Brun's avatar
Sophie Brun committed
30 31 32 33 34
		clients: make(map[string]*Station),
	}
}

func (ap *AccessPoint) MarshalJSON() ([]byte, error) {
Sophie Brun's avatar
Sophie Brun committed
35 36 37
	ap.Lock()
	defer ap.Unlock()

Sophie Brun's avatar
Sophie Brun committed
38
	doc := apJSON{
Sophie Brun's avatar
Sophie Brun committed
39 40 41
		Station:   ap.Station,
		Clients:   make([]*Station, 0),
		Handshake: ap.withKeyMaterial,
Sophie Brun's avatar
Sophie Brun committed
42 43 44 45 46 47 48 49 50 51 52 53 54 55
	}

	for _, c := range ap.clients {
		doc.Clients = append(doc.Clients, c)
	}

	return json.Marshal(doc)
}

func (ap *AccessPoint) Get(bssid string) (*Station, bool) {
	ap.Lock()
	defer ap.Unlock()

	bssid = NormalizeMac(bssid)
Sophie Brun's avatar
Sophie Brun committed
56
	if s, found := ap.clients[bssid]; found {
Sophie Brun's avatar
Sophie Brun committed
57 58 59 60 61
		return s, true
	}
	return nil, false
}

Sophie Brun's avatar
Sophie Brun committed
62 63 64 65 66 67 68 69 70 71
func (ap *AccessPoint) RemoveClient(mac string) {
	ap.Lock()
	defer ap.Unlock()

	bssid := NormalizeMac(mac)
	if _, found := ap.clients[bssid]; found {
		delete(ap.clients, bssid)
	}
}

Sophie Brun's avatar
Sophie Brun committed
72
func (ap *AccessPoint) AddClientIfNew(bssid string, frequency int, rssi int8) (*Station, bool) {
Sophie Brun's avatar
Sophie Brun committed
73 74 75 76
	ap.Lock()
	defer ap.Unlock()

	bssid = NormalizeMac(bssid)
Sophie Brun's avatar
Sophie Brun committed
77
	alias := ap.aliases.GetOr(bssid, "")
Sophie Brun's avatar
Sophie Brun committed
78

Sophie Brun's avatar
Sophie Brun committed
79
	if s, found := ap.clients[bssid]; found {
Sophie Brun's avatar
Sophie Brun committed
80 81 82 83 84
		// update
		s.Frequency = frequency
		s.RSSI = rssi
		s.LastSeen = time.Now()

Sophie Brun's avatar
Sophie Brun committed
85 86 87 88
		if alias != "" {
			s.Alias = alias
		}

Sophie Brun's avatar
Sophie Brun committed
89
		return s, false
Sophie Brun's avatar
Sophie Brun committed
90 91 92
	}

	s := NewStation("", bssid, frequency, rssi)
Sophie Brun's avatar
Sophie Brun committed
93
	s.Alias = alias
Sophie Brun's avatar
Sophie Brun committed
94 95
	ap.clients[bssid] = s

Sophie Brun's avatar
Sophie Brun committed
96
	return s, true
Sophie Brun's avatar
Sophie Brun committed
97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
}

func (ap *AccessPoint) NumClients() int {
	ap.Lock()
	defer ap.Unlock()
	return len(ap.clients)
}

func (ap *AccessPoint) Clients() (list []*Station) {
	ap.Lock()
	defer ap.Unlock()

	list = make([]*Station, 0)
	for _, c := range ap.clients {
		list = append(list, c)
	}
	return
}
Sophie Brun's avatar
Sophie Brun committed
115 116 117 118 119 120 121 122 123 124

func (ap *AccessPoint) EachClient(cb func(mac string, station *Station)) {
	ap.Lock()
	defer ap.Unlock()

	for m, station := range ap.clients {
		cb(m, station)
	}
}

Sophie Brun's avatar
Sophie Brun committed
125 126 127 128 129 130 131 132 133 134 135 136 137 138
func (ap *AccessPoint) WithKeyMaterial(state bool) {
	ap.Lock()
	defer ap.Unlock()

	ap.withKeyMaterial = state
}

func (ap *AccessPoint) HasKeyMaterial() bool {
	ap.Lock()
	defer ap.Unlock()

	return ap.withKeyMaterial
}

Sophie Brun's avatar
Sophie Brun committed
139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156
func (ap *AccessPoint) NumHandshakes() int {
	ap.Lock()
	defer ap.Unlock()

	sum := 0

	for _, c := range ap.clients {
		if c.Handshake.Complete() {
			sum++
		}
	}

	return sum
}

func (ap *AccessPoint) HasHandshakes() bool {
	return ap.NumHandshakes() > 0
}
Sophie Brun's avatar
Sophie Brun committed
157 158 159 160 161 162 163 164 165 166 167 168 169

func (ap *AccessPoint) HasPMKID() bool {
	ap.Lock()
	defer ap.Unlock()

	for _, c := range ap.clients {
		if c.Handshake.HasPMKID() {
			return true
		}
	}

	return false
}