From a8bfd8f183ec91b2651605dc96014993b523e253 Mon Sep 17 00:00:00 2001
From: root <root@game.hexor.ru>
Date: Thu, 6 May 2021 18:23:36 +0300
Subject: [PATCH] Fix bugs

---
 bot.py |  2 +-
 gen.py | 20 ++++++++++++++------
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/bot.py b/bot.py
index 5cc6c09..63e87bb 100644
--- a/bot.py
+++ b/bot.py
@@ -26,11 +26,11 @@ def add_peer(update, context):
         return False
     if len(update.message.text.split()) < 2:
         update.message.reply_text('Wrong usage.\n<b>Help:</b>\n <b>*</b> /add <i>peer name</i>', parse_mode='HTML', disable_web_page_preview=True)
+        return False
     log.info(update.message.chat.username)
     peer_name = "_".join(update.message.text.split()[1:])
     log.info("Creating peer %s", peer_name)
     wg_add_peer(peer_name)
-    update_configs()
     update.message.reply_photo(open(f'clients/{peer_name}-qr.png', 'rb'), filename=f'{peer_name} QR.png', quote=True, caption=open(f'clients/{peer_name}.conf', 'r').read())
 
 
diff --git a/gen.py b/gen.py
index 218f9cd..1860711 100644
--- a/gen.py
+++ b/gen.py
@@ -39,13 +39,13 @@ class Peer:
                 else:
                     self.priv_key = data[0].split(':')[1].strip()
                     self.comment = data[1].split(':')[1].strip()
-                    self.managed = True
+#                   self.managed = True
             except:
                 pass
         else:
             self.priv_key = self.generate_key()
             self.pub_key = self.public_key(self.priv_key)
-            self.allowed_ips = allowed_ips
+            self.allowed_ips = allowed_ips if allowed_ips else Helper().next_ip
 
         self.full_comment = "# priv_key: " + " ; ".join([self.priv_key, "comment: " + comment])
 
@@ -68,7 +68,6 @@ class Peer:
 #           return False
         filename = f"clients/{self.comment.replace(' ', '_')}"
         _wg = wgconfig.WGConfig(f"{filename}.conf")
-        #_wg.initialize_file(f"# Autogenerated file \n# Comment: {self.comment}\n# Public key: {self.pub_key}\n") 
         _wg.initialize_file() 
         _wg.add_attr(None, 'Address', self.allowed_ips)
         _wg.add_attr(None, 'DNS', helper.dns)
@@ -97,10 +96,12 @@ class Helper:
 
     @property
     def server_pub_key(self):
+        """Return server public key"""
         return Peer().public_key(self.wg.interface['PrivateKey'])
 
     @property
     def peer_list(self):
+        """Return list of WG peers"""
         peer_list = list()
         for i, v in self.wg.peers.items():
             peer_list.append(Peer(peer=v))
@@ -108,6 +109,7 @@ class Helper:
 
     @property
     def ip_list(self):
+        """Return list of IPs"""
         ip_list = list()
         ip_list.append(ipaddress.ip_address(Helper().wg.interface['Address'].split('/')[0]))
         for i, v in self.wg.peers.items():
@@ -122,9 +124,15 @@ class Helper:
                 pass
         ip_list.sort()
         return ip_list
+    
+    @property
+    def next_ip(self):
+        """Return next free IP"""
+        return self.ip_list[-1]+1
 
 
-    def add_client(self, comment):
+    def add_peer(self, comment):
+        """Generate a new peer"""
         cl = Peer(comment=comment)
         self.wg.add_peer(cl.pub_key, cl.full_comment)
         self.wg.add_attr(cl.pub_key, 'AllowedIPs', f"{self.ip_list[-1]+1}/32")
@@ -133,15 +141,15 @@ class Helper:
 def add_peer(peer_name):
     print('Generate a new peer config.')
     helper = Helper()
-    helper.add_client(peer_name)
+    helper.add_peer(peer_name)
     helper.wg.write_file()
+    system('systemctl restart wg-quick@wg0.service')
 
 def update_configs():
     print("Update all clients configs.")
     helper = Helper()
     for peer in helper.peer_list:
         peer.gen_config(Helper())
-    system('systemctl restart wg-quick@wg0.service')
 
 if not is_update and peer_name:
     add_peer(peer_name)