From e9b93b8586700729bd0d11abfe771efd23045919 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Wed, 6 Feb 2019 14:07:58 +0100
Subject: [PATCH] fix: save unsaved packets only if handshake is complete

---
 network/wifi.go | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/network/wifi.go b/network/wifi.go
index 75dcff2a..4a50e167 100644
--- a/network/wifi.go
+++ b/network/wifi.go
@@ -220,11 +220,16 @@ func (w *WiFi) SaveHandshakesTo(fileName string, linkType layers.LinkType) error
 
 	for _, ap := range w.aps {
 		for _, station := range ap.Clients() {
-			station.Handshake.EachUnsavedPacket(func(pkt gopacket.Packet) {
-				err = writer.WritePacket(pkt.Metadata().CaptureInfo, pkt.Data())
-			})
-			if err != nil {
-				return err
+			if station.Handshake.Complete() {
+				err = nil
+				station.Handshake.EachUnsavedPacket(func(pkt gopacket.Packet) {
+					if err == nil {
+						err = writer.WritePacket(pkt.Metadata().CaptureInfo, pkt.Data())
+					}
+				})
+				if err != nil {
+					return err
+				}
 			}
 		}
 	}