From 268d26d802dfac4b825cb7aa72d10ae6d63adbec Mon Sep 17 00:00:00 2001
From: Jake Hillion <jake@hillion.co.uk>
Date: Mon, 9 Nov 2020 10:54:40 +0000
Subject: [PATCH] adding support for client behind

---
 structure/structure.py | 72 ++++++++++++++++++++++++++++--------------
 1 file changed, 48 insertions(+), 24 deletions(-)

diff --git a/structure/structure.py b/structure/structure.py
index 335e05c..d2e8614 100644
--- a/structure/structure.py
+++ b/structure/structure.py
@@ -126,15 +126,19 @@ class Node:
 
 
 class SpeedTestServer(Node):
-    def client(self, server: Interface):
-        pass
+    def get_internet_setup(self) -> Optional[str]:
+        return textwrap.dedent('''
+            cloud-init status --wait || cloud-init status --long
+            sudo apt-get install -y iperf3
+        ''')
 
-    # Entry method for running the serve with `with speedtest:`
-    def __enter__(self):
-        pass
+    def server(self):
+        self.ssh('iperf3 -s -1 -D', error_stdout=True, error_stderr=True)
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        pass
+    def client(self, target, time=30):
+        command = 'iperf3 -c {target} -t {time} -O 5 -J'.format(target=target, time=time)
+        out = self.ssh(command, error_stdout=True, error_stderr=True, return_stdout=True)
+        return json.loads(out)
 
 
 class RemotePortal(Node):
@@ -166,28 +170,34 @@ class RemotePortal(Node):
     def get_setup(self) -> Optional[str]:
         return textwrap.dedent('''
             set -e
-            
-            sudo sysctl -w net.ipv4.conf.all.arp_announce=1
-            sudo sysctl -w net.ipv4.conf.all.arp_ignore=2
-                        
+
             cat << EOF > config.ini
             [Host]
             PrivateKey = INVALID
-            
+
             [Peer]
             PublicKey = INVALID
             Method = TCP
-            
+
             LocalHost = {local_host}
             LocalPort = 1234
             EOF
-            
+
             (nohup sudo ./mpbl3p > mpbl3p.log 2>&1 & echo $! > mpbl3p.pid)
-            
+
             sleep 1
             sudo ip link set up nc0
             sudo ip addr add 172.19.152.2/31 dev nc0
             
+            ip rule add from all table local priority 20
+            ip rule del 0
+
+            ip rule add to {local_host} dport 1234 table local priority 9
+
+            ip route flush 10
+            ip route add table 10 to {local_host} via 172.19.152.3 dev nc0
+            ip rule add to {local_host} table 10 priority 10
+
             ps $(cat mpbl3p.pid)
         ''').format(
             local_host=self.get_interfaces()[0].get_address(),
@@ -239,9 +249,9 @@ class LocalPortal(Node):
             [Peer]
             PublicKey = INVALID
             Method = TCP
-            
+
             LocalHost = {local_host}
-            
+
             RemoteHost = {remote_host}
             RemotePort = 1234
         ''')
@@ -266,10 +276,10 @@ class LocalPortal(Node):
 
         return textwrap.dedent('''
             set -e
-            
+
             sudo sysctl -w net.ipv4.conf.all.arp_announce=1
-            sudo sysctl -w net.ipv4.conf.all.arp_ignore=2
-                        
+            sudo sysctl -w net.ipv4.conf.all.arp_ignore=1
+
             {policy_routing}
 
             cat << EOF > config.ini
@@ -278,15 +288,29 @@ class LocalPortal(Node):
 
             {peers}
             EOF
-            
+
             (nohup sudo ./mpbl3p > mpbl3p.log 2>&1 & echo $! > mpbl3p.pid)
-            
+
             sleep 1
             sudo ip link set up nc0
             sudo ip addr add 172.19.152.3/31 dev nc0
-            
+
+            ip route flush 8
+            ip route add table 8 default via 172.19.152.2 dev nc0
+            ip rule add from {remote_host} iif {local_interface} table 8 priority 8
+
+            ip route flush 9
+            ip route add table 9 to {remote_host} dev {local_interface}
+            ip rule add to {remote_host} table 9 priority 9
+
             ps $(cat mpbl3p.pid)
-        ''').format(**self.setup_params, peers=peers, policy_routing=policy_routing)
+        ''').format(
+            **self.setup_params,
+            peers=peers,
+            policy_routing=policy_routing,
+            remote_host=self.remote_portal.get_interfaces()[0].get_address(),
+            local_interface='eth{}'.format(len(self.get_interfaces())-2),
+        )
 
     def speedtest_server(self):
         self.ssh('iperf3 -s -1 -D', error_stdout=True, error_stderr=True)