From 8907ed33f86c59aa8a1b161f7a351d22891890bb Mon Sep 17 00:00:00 2001
From: Fedor Lyanguzov <Fedor-Lyanguzov@users.noreply.github.com>
Date: Sat, 28 Dec 2024 01:09:24 +0300
Subject: [PATCH] add interface section formatting

---
 .activate               |  7 +++++++
 tests/test_app.py       | 15 ++++++++++++--
 vpn_manager/__init__.py |  1 -
 vpn_manager/__main__.py | 46 ++++++++++++++++++++++++++++++++++++-----
 4 files changed, 61 insertions(+), 8 deletions(-)
 create mode 100644 .activate

diff --git a/.activate b/.activate
new file mode 100644
index 0000000..a3607ff
--- /dev/null
+++ b/.activate
@@ -0,0 +1,7 @@
+
+alias venv='source .venv/Scripts/activate'
+alias check='flake8 vpn_manager'
+alias format='black vpn_manager'
+alias test='pytest'
+
+venv
diff --git a/tests/test_app.py b/tests/test_app.py
index b089239..6e7523d 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -1,6 +1,17 @@
 from vpn_manager import *
 
 def test_format_static_peer():
-    sp = StaticPeer('sample-public-key', '127.0.0.1', '12345')
+    sp = Peer('PUB', "PRV", '10.0.0.1/32', '12345', '127.0.0.1')
     peer_section = format_static_peer(sp, '0.0.0.0/0')
-    assert peer_section == '[Peer]\nPublicKey = sample-public-key\nAllowedIPs = 0.0.0.0/0\nEndpoint = 127.0.0.1:12345\nPersistentKeepAlive = 30'
+    assert peer_section == '[Peer]\nPublicKey = PUB\nAllowedIPs = 0.0.0.0/0\nEndpoint = 127.0.0.1:12345\nPersistentKeepAlive = 30'
+
+
+def test_true():
+    assert format_interface(Peer('PUB', 'PRV', '10.0.0.1/32', 'PORT'), '1.1.1.1') \
+        == '[Interface]\nPrivateKey = PRV\nAddress = 10.0.0.1/32\nListenPort = PORT\nDNS = 1.1.1.1'
+    assert format_interface(Peer('PUB', 'PRV', '10.0.0.1/32', None), '1.1.1.1') \
+        == '[Interface]\nPrivateKey = PRV\nAddress = 10.0.0.1/32\nDNS = 1.1.1.1'
+    assert format_interface(Peer('PUB', 'PRV', '10.0.0.1/32', 'PORT'), '1.1.1.1', forward=True) \
+        == '[Interface]\nPrivateKey = PRV\nAddress = 10.0.0.1/32\nListenPort = PORT\nDNS = 1.1.1.1\nPostUp = iptables -A FORWARD -i wg0 -j ACCEPT; iptables -t nat -A POSTROUTING -o eth0 -j MASQUERADE\nPostDown = iptables -D FORWARD -i wg0 -j ACCEPT; iptables -t nat -D POSTROUTING -o eth0 -j MASQUERADE'
+    assert format_interface(Peer('PUB', 'PRV', '10.0.0.1/32', None), '1.1.1.1', forward=True) \
+        == '[Interface]\nPrivateKey = PRV\nAddress = 10.0.0.1/32\nDNS = 1.1.1.1\nPostUp = iptables -A FORWARD -i wg0 -j ACCEPT; iptables -t nat -A POSTROUTING -o eth0 -j MASQUERADE\nPostDown = iptables -D FORWARD -i wg0 -j ACCEPT; iptables -t nat -D POSTROUTING -o eth0 -j MASQUERADE'
diff --git a/vpn_manager/__init__.py b/vpn_manager/__init__.py
index 737754e..a86c6bb 100644
--- a/vpn_manager/__init__.py
+++ b/vpn_manager/__init__.py
@@ -1,2 +1 @@
 from .__main__ import *
-
diff --git a/vpn_manager/__main__.py b/vpn_manager/__main__.py
index 46b7f94..9470d5b 100644
--- a/vpn_manager/__main__.py
+++ b/vpn_manager/__main__.py
@@ -1,19 +1,55 @@
 from dataclasses import dataclass
+from textwrap import dedent
+
 
 @dataclass
-class StaticPeer:
+class Peer:
     public_key: str
-    endpoint: str
+    private_key: str
+    address_cidr: str
     port: str
+    endpoint: str = None
 
 
 def format_static_peer(static_peer, routes, keepalive=30):
-    return f'''[Peer]
+    return f"""
+[Peer]
 PublicKey = {static_peer.public_key}
 AllowedIPs = {routes}
 Endpoint = {static_peer.endpoint}:{static_peer.port}
-PersistentKeepAlive = {keepalive}'''
-    
+PersistentKeepAlive = {keepalive}
+""".strip()
+
+
+def format_interface(peer, dns, forward=False):
+    if forward:
+        forward = dedent(
+            """\
+            PostUp = iptables -A FORWARD -i wg0 -j ACCEPT; iptables -t nat -A POSTROUTING -o eth0 -j MASQUERADE
+            PostDown = iptables -D FORWARD -i wg0 -j ACCEPT; iptables -t nat -D POSTROUTING -o eth0 -j MASQUERADE
+            """
+        )
+        # wg0 always?
+        # eth0 always?
+    else:
+        forward = ""
+    if peer.port:
+        port = f"ListenPort = {peer.port}\n"
+    else:
+        port = ""
+    dns = f"DNS = {dns}\n"
+    return f"""
+[Interface]
+PrivateKey = {peer.private_key}
+Address = {peer.address_cidr}
+{port}\
+{dns}\
+{forward}\
+""".strip()
+
+
+def remove_empty_lines(text):
+    return "\n".join([s for s in text.splitlines() if s.strip()])
 
 
 def main():