diff --git a/src/helper/wifiStateAnalyzer.js b/src/helper/wifiStateAnalyzer.js
index bb71f60..e2622de 100644
--- a/src/helper/wifiStateAnalyzer.js
+++ b/src/helper/wifiStateAnalyzer.js
@@ -1,3 +1,4 @@
+const { HandshakeStage } = require.main.require('./dto/Packet.js');
 
 function keyInfoFromRaw(keyInfoRaw) {
     return {
@@ -15,8 +16,28 @@ function keyInfoFromRaw(keyInfoRaw) {
     };
 }
 
+const HANDSHAKE_STAGE_KEYINFO = {
+    "keys": ["Install", "KeyACK", "KeyMIC", "Secure"],
+    "0100": HandshakeStage[1],
+    "0010": HandshakeStage[2],
+    "1111": HandshakeStage[3],
+    "0011": HandshakeStage[4],
+};
+function handshakeStageFromKeyInfo(keyInfo){
+
+    // Extract compare-keys
+    let keyData = "";
+    for (const key of HANDSHAKE_STAGE_KEYINFO['keys']) {
+        keyData += keyInfo[key].toString();
+    }
+    
+    // Get and return stage
+    return  HANDSHAKE_STAGE_KEYINFO[keyData];
+}
+
 
 // Specify exports
 module.exports = {
     keyInfoFromRaw,
+    handshakeStageFromKeyInfo,
 };
\ No newline at end of file
diff --git a/src/streamHandler/PacketStreamFactory.js b/src/streamHandler/PacketStreamFactory.js
index 3dbbe26..9d7c866 100644
--- a/src/streamHandler/PacketStreamFactory.js
+++ b/src/streamHandler/PacketStreamFactory.js
@@ -135,13 +135,9 @@ class PacketStreamFactory extends Transform{
 
                 // Read key-information
                 const keyInfoRaw = (packet.payloadData[0x5]<<0x8) + packet.payloadData[0x6];
-                const keyInfo = wifiStateAnalyser.keyInfoFromRaw(keyInfoRaw);
-                
-                newPacket.handshakeStage =  (!keyInfo.Install &&  keyInfo.KeyACK && !keyInfo.KeyMIC && !keyInfo.Secure)? HandshakeStage[1] :
-                                            (!keyInfo.Install && !keyInfo.KeyACK &&  keyInfo.KeyMIC && !keyInfo.Secure)? HandshakeStage[2] :
-                                            ( keyInfo.Install &&  keyInfo.KeyACK &&  keyInfo.KeyMIC &&  keyInfo.Secure)? HandshakeStage[3] :
-                                            (!keyInfo.Install && !keyInfo.KeyACK &&  keyInfo.KeyMIC &&  keyInfo.Secure)? HandshakeStage[4] :
-                                            null;
+                const keyInfo = wifiStateAnalyser.keyInfoFromRaw(keyInfoRaw);   // Convert
+
+                newPacket.handshakeStage =  wifiStateAnalyser.handshakeStageFromKeyInfo(keyInfo);   // Get stage
                 break;
         }
         if(newPacket) packet = Object.assign(newPacket, packet);