From 4ea9ab22f6725ada329341518210486499e7dbf3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89milio=20Gonzalez?= <egg997@gmail.com>
Date: Sat, 8 Jun 2019 17:45:04 -0700
Subject: [PATCH] Add a connection time log statement upon client disconnection

---
 pyrdp/mitm/TCPMITM.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/pyrdp/mitm/TCPMITM.py b/pyrdp/mitm/TCPMITM.py
index 7feaac52b..4b08162d4 100644
--- a/pyrdp/mitm/TCPMITM.py
+++ b/pyrdp/mitm/TCPMITM.py
@@ -3,7 +3,7 @@
 # Copyright (C) 2019 GoSecure Inc.
 # Licensed under the GPLv3 or later.
 #
-
+import time
 from logging import LoggerAdapter
 from typing import Coroutine
 
@@ -28,6 +28,9 @@ def __init__(self, client: TwistedTCPLayer, server: TwistedTCPLayer, attacker: T
         :param serverConnector: coroutine that connects to the server side, closed when the client disconnects
         """
 
+        self.connectionTime = 0
+        # To keep track of the duration of the TCP connection.
+
         self.client = client
         self.server = server
         self.attacker = attacker
@@ -63,7 +66,7 @@ def onClientConnection(self):
         """
         Log the fact that a new client has connected.
         """
-
+        self.connectionTime = time.time()
         self.log.info("New client connected")
 
     def onClientDisconnection(self, reason):
@@ -72,8 +75,11 @@ def onClientDisconnection(self, reason):
         :param reason: reason for disconnection
         """
 
+        self.connectionTime = time.time() - self.connectionTime
+
         self.recordConnectionClose()
         self.log.info("Client connection closed. %(reason)s", {"reason": reason.value})
+        self.log.info("Client connection time: %(connectionTime)s secs", {"connectionTime": self.connectionTime})
         self.serverConnector.close()
         self.server.disconnect(True)