wifi_handshake.go 2.12 KB
Newer Older
Sophie Brun's avatar
Sophie Brun committed
1 2 3 4
package network

import (
	"sync"
Sophie Brun's avatar
Sophie Brun committed
5 6 7

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
Sophie Brun's avatar
Sophie Brun committed
8 9 10 11 12 13 14 15 16
)

type Handshake struct {
	sync.Mutex

	Beacon        gopacket.Packet
	Challenges    []gopacket.Packet
	Responses     []gopacket.Packet
	Confirmations []gopacket.Packet
Sophie Brun's avatar
Sophie Brun committed
17
	hasPMKID      bool
Sophie Brun's avatar
Sophie Brun committed
18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39
	unsaved       []gopacket.Packet
}

func NewHandshake() *Handshake {
	return &Handshake{
		Challenges:    make([]gopacket.Packet, 0),
		Responses:     make([]gopacket.Packet, 0),
		Confirmations: make([]gopacket.Packet, 0),
		unsaved:       make([]gopacket.Packet, 0),
	}
}

func (h *Handshake) SetBeacon(pkt gopacket.Packet) {
	h.Lock()
	defer h.Unlock()

	if h.Beacon == nil {
		h.Beacon = pkt
		h.unsaved = append(h.unsaved, pkt)
	}
}

Sophie Brun's avatar
Sophie Brun committed
40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65
func (h *Handshake) AddAndGetPMKID(pkt gopacket.Packet) []byte {
	h.AddFrame(0, pkt)

	prevWasKey := false
	for _, layer := range pkt.Layers() {
		if layer.LayerType() == layers.LayerTypeEAPOLKey {
			prevWasKey = true
			continue
		}

		if prevWasKey && layer.LayerType() == layers.LayerTypeDot11InformationElement {
			info := layer.(*layers.Dot11InformationElement)
			if info.ID == layers.Dot11InformationElementIDVendor && info.Length == 20 {
				h.Lock()
				defer h.Unlock()
				h.hasPMKID = true
				return info.Info
			}
		}

		prevWasKey = false
	}

	return nil
}

Sophie Brun's avatar
Sophie Brun committed
66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92
func (h *Handshake) AddFrame(n int, pkt gopacket.Packet) {
	h.Lock()
	defer h.Unlock()

	switch n {
	case 0:
		h.Challenges = append(h.Challenges, pkt)
	case 1:
		h.Responses = append(h.Responses, pkt)
	case 2:
		h.Confirmations = append(h.Confirmations, pkt)
	}

	h.unsaved = append(h.unsaved, pkt)
}

func (h *Handshake) Complete() bool {
	h.Lock()
	defer h.Unlock()

	nChal := len(h.Challenges)
	nResp := len(h.Responses)
	nConf := len(h.Confirmations)

	return nChal > 0 && nResp > 0 && nConf > 0
}

Sophie Brun's avatar
Sophie Brun committed
93 94 95 96 97 98
func (h *Handshake) HasPMKID() bool {
	h.Lock()
	defer h.Unlock()
	return h.hasPMKID
}

Sophie Brun's avatar
Sophie Brun committed
99 100 101 102 103 104 105 106 107 108 109 110 111 112 113
func (h *Handshake) NumUnsaved() int {
	h.Lock()
	defer h.Unlock()
	return len(h.unsaved)
}

func (h *Handshake) EachUnsavedPacket(cb func(gopacket.Packet)) {
	h.Lock()
	defer h.Unlock()

	for _, pkt := range h.unsaved {
		cb(pkt)
	}
	h.unsaved = make([]gopacket.Packet, 0)
}