From 2e32f188d086c629bcfb343d0c6bf1375e4e8821 Mon Sep 17 00:00:00 2001
From: Lukas Kupczyk <lukas.kupczyk@gmail.com>
Date: Mon, 17 Oct 2022 13:50:34 +0200
Subject: [PATCH] Allow to specify listen address for pyrdp-mitm

---
 bin/pyrdp-mitm.py    | 6 +++---
 pyrdp/mitm/cli.py    | 2 ++
 pyrdp/mitm/config.py | 3 +++
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/bin/pyrdp-mitm.py b/bin/pyrdp-mitm.py
index 3ae5c43dd..f1e843e2c 100755
--- a/bin/pyrdp-mitm.py
+++ b/bin/pyrdp-mitm.py
@@ -37,13 +37,13 @@ def main():
         except Exception:
             logger.warning('Unable to set transparent socket. Are you running as root?')
 
-    s.bind(('0.0.0.0', config.listenPort))
+    s.bind((config.listenAddress, config.listenPort))
     s.listen()  # Non-blocking.
     reactor.adoptStreamPort(s.fileno(), socket.AF_INET, MITMServerFactory(config))
     s.close()  # reactor creates a copy of the fd.
 
-    message = "MITM Server listening on 0.0.0.0:%(port)d"
-    params = {"port": config.listenPort}
+    message = "MITM Server listening on %(address)s:%(port)d"
+    params = {"address": config.listenAddress, "port": config.listenPort}
 
     if "HOST_IP" in os.environ:
         message += ". Host IP: %(host_ip)s"
diff --git a/pyrdp/mitm/cli.py b/pyrdp/mitm/cli.py
index 94d9938f6..2674acd32 100644
--- a/pyrdp/mitm/cli.py
+++ b/pyrdp/mitm/cli.py
@@ -76,6 +76,7 @@ def showConfiguration(config: MITMConfig):
 def buildArgParser():
     parser = argparse.ArgumentParser()
     parser.add_argument("target", help="IP:port of the target RDP machine (ex: 192.168.1.10:3390)", nargs='?', default=None)
+    parser.add_argument("-a", "--address", help="Address to listen on (default: 0.0.0.0)", default="0.0.0.0")
     parser.add_argument("-l", "--listen", help="Port number to listen on (default: 3389)", default=3389)
     parser.add_argument("-o", "--output", help="Output folder", default="pyrdp_output")
     parser.add_argument("-i", "--destination-ip",
@@ -181,6 +182,7 @@ def configure(cmdline=None) -> MITMConfig:
     config.targetHost = targetHost
     config.targetPort = targetPort
     config.privateKeyFileName = key
+    config.listenAddress = args.address
     config.listenPort = int(args.listen)
     config.certificateFileName = certificate
     config.attackerHost = args.destination_ip
diff --git a/pyrdp/mitm/config.py b/pyrdp/mitm/config.py
index f0946a852..47a143b65 100644
--- a/pyrdp/mitm/config.py
+++ b/pyrdp/mitm/config.py
@@ -22,6 +22,9 @@ def __init__(self):
         self.targetPort: int = None
         """The RDP server's port"""
 
+        self.listenAddress: str = "0.0.0.0"
+        """The address to listen on."""
+
         self.listenPort: int = 3389
         """The port to bind for listening."""