From 559a51393b38c7f8c2d0a7dcaa8e5927c5cad51d Mon Sep 17 00:00:00 2001
From: k4yt3x <k4yt3x@k4yt3x.com>
Date: Fri, 21 May 2021 21:37:28 +0000
Subject: [PATCH] added support for optional peer attributes, added __main__

---
 wg_meshconf/__main__.py         | 14 ++++++++++++++
 wg_meshconf/database_manager.py | 12 ++++++++++++
 wg_meshconf/wg_meshconf.py      | 13 ++++++-------
 3 files changed, 32 insertions(+), 7 deletions(-)
 create mode 100644 wg_meshconf/__main__.py

diff --git a/wg_meshconf/__main__.py b/wg_meshconf/__main__.py
new file mode 100644
index 0000000..ccabdca
--- /dev/null
+++ b/wg_meshconf/__main__.py
@@ -0,0 +1,14 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+Name: wg-meshconf __main__
+Creator: K4YT3X
+Date Created: May 21, 2021
+Last Modified: May 21, 2021
+"""
+
+from .wg_meshconf import main
+
+# launch the main function if it is not imported as a package
+if __name__ == "__main__":
+    main()
diff --git a/wg_meshconf/database_manager.py b/wg_meshconf/database_manager.py
index 4f033a4..7e4d5d7 100755
--- a/wg_meshconf/database_manager.py
+++ b/wg_meshconf/database_manager.py
@@ -60,6 +60,10 @@ PEER_ATTRIBUTES = [
     "PersistentKeepalive",
 ]
 
+PEER_OPTIONAL_ATTRIBUTES = [
+    "PersistentKeepalive",
+]
+
 
 class DatabaseManager:
     def __init__(self, database_path: pathlib.Path):
@@ -95,6 +99,7 @@ class DatabaseManager:
         Endpoint: str = None,
         AllowedIPs: list = None,
         ListenPort: int = None,
+        PersistentKeepalive: int = None,
         FwMark: str = None,
         PrivateKey: str = None,
         DNS: str = None,
@@ -133,6 +138,7 @@ class DatabaseManager:
         Endpoint: str = None,
         AllowedIPs: list = None,
         ListenPort: int = None,
+        PersistentKeepalive: int = None,
         FwMark: str = None,
         PrivateKey: str = None,
         DNS: str = None,
@@ -303,3 +309,9 @@ class DatabaseManager:
                         else:
                             allowed_ips = ", ".join(database["peers"][p]["Address"])
                         config.write("AllowedIPs = {}\n".format(allowed_ips))
+
+                    for key in PEER_OPTIONAL_ATTRIBUTES:
+                        if database["peers"][p].get(key) is not None:
+                            config.write(
+                                "{} = {}\n".format(key, database["peers"][p][key])
+                            )
diff --git a/wg_meshconf/wg_meshconf.py b/wg_meshconf/wg_meshconf.py
index ef7f7d7..2d5eec8 100755
--- a/wg_meshconf/wg_meshconf.py
+++ b/wg_meshconf/wg_meshconf.py
@@ -4,7 +4,7 @@
 Name: wg-meshconf
 Creator: K4YT3X
 Date Created: July 19, 2020
-Last Modified: January 12, 2021
+Last Modified: May 21, 2021
 
 Licensed under the GNU General Public License Version 3 (GNU GPL v3),
     available at: https://www.gnu.org/licenses/gpl-3.0.txt
@@ -52,6 +52,7 @@ def parse_arguments():
     )
     addpeer.add_argument("--privatekey", help="private key of server interface")
     addpeer.add_argument("--listenport", help="port to listen on", default=51820)
+    addpeer.add_argument("--persistentkeepalive", help="set persistent keepalive interval")
     addpeer.add_argument("--fwmark", help="fwmark for outgoing packets")
     addpeer.add_argument("--dns", help="server interface DNS servers")
     addpeer.add_argument("--mtu", help="server interface MTU")
@@ -77,6 +78,7 @@ def parse_arguments():
     )
     updatepeer.add_argument("--privatekey", help="private key of server interface")
     updatepeer.add_argument("--listenport", help="port to listen on")
+    updatepeer.add_argument("--persistentkeepalive", help="set persistent keepalive interval")
     updatepeer.add_argument("--fwmark", help="fwmark for outgoing packets")
     updatepeer.add_argument("--dns", help="server interface DNS servers")
     updatepeer.add_argument("--mtu", help="server interface MTU")
@@ -129,7 +131,7 @@ def parse_arguments():
         "--output",
         help="configuration file output directory",
         type=pathlib.Path,
-        default=pathlib.Path(__file__).parent.absolute() / "output",
+        default=pathlib.Path.cwd() / "output",
     )
 
     return parser.parse_args()
@@ -149,6 +151,7 @@ def main():
             args.endpoint,
             args.allowedips,
             args.listenport,
+            args.persistentkeepalive,
             args.fwmark,
             args.privatekey,
             args.dns,
@@ -168,6 +171,7 @@ def main():
             args.endpoint,
             args.allowedips,
             args.listenport,
+            args.persistentkeepalive,
             args.fwmark,
             args.privatekey,
             args.dns,
@@ -195,8 +199,3 @@ def main():
             "No command specified\nUse wg-meshconf --help to see available commands",
             file=sys.stderr,
         )
-
-
-# launch the main function if it is not imported as a package
-if __name__ == "__main__":
-    main()