From 09ecf58753e483b633e5a92980f401e001a444f3 Mon Sep 17 00:00:00 2001
From: k4yt3x <k4yt3x@protonmail.com>
Date: Wed, 17 Oct 2018 21:06:37 -0400
Subject: [PATCH] 1.1.4 serialized profile storing

---
 wireguard_mesh_configurator.py | 67 ++++++++++++++++------------------
 1 file changed, 31 insertions(+), 36 deletions(-)

diff --git a/wireguard_mesh_configurator.py b/wireguard_mesh_configurator.py
index 79799b4..479aace 100755
--- a/wireguard_mesh_configurator.py
+++ b/wireguard_mesh_configurator.py
@@ -11,15 +11,15 @@ Licensed under the GNU General Public License Version 3 (GNU GPL v3),
 (C) 2018 K4YT3X
 """
 import avalon_framework as avalon
-import json
 import os
+import pickle
 import re
 import readline
 import subprocess
 import sys
 import traceback
 
-VERSION = '1.1.3'
+VERSION = '1.1.4'
 COMMANDS = [
     'Interactive',
     'ShowPeers',
@@ -33,6 +33,18 @@ COMMANDS = [
 ]
 
 
+class Utilities:
+    """ Useful utilities
+
+    This class contains a number of utility tools.
+    """
+
+    def execute(command, input_value=''):
+        process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
+        output = process.communicate(input=input_value)[0]
+        return output.decode().replace('\n', '')
+
+
 class ShellCompleter(object):
     """ A Cisco-IOS-like shell completer
 
@@ -93,8 +105,7 @@ class WireGuard:
         Generate a new wireguard private key via
         wg command.
         """
-        output = subprocess.Popen(['wg', 'genkey'], stdout=subprocess.PIPE).communicate()[0]
-        return output.decode().replace('\n', '')
+        return Utilities.execute(['wg', 'genkey'])
 
     def pubkey(self, public_key):
         """ Convert WG private key into public key
@@ -102,15 +113,12 @@ class WireGuard:
         Uses wg pubkey command to convert the wg private
         key into a public key.
         """
-        process = subprocess.Popen(['wg', 'pubkey'], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
-        output = process.communicate(input=public_key.encode('utf-8'))[0]
-        return output.decode().replace('\n', '')
+        return Utilities.execute(['wg', 'pubkey'], input_value=public_key.encode('utf-8'))
 
     def genpsk(self):
         """ Generate a random base64 psk
         """
-        output = subprocess.Popen(['wg', 'genpsk'], stdout=subprocess.PIPE).communicate()[0]
-        return output.decode().replace('\n', '')
+        return Utilities.execute(['wg', 'genpsk'])
 
 
 class ProfileManager(object):
@@ -126,35 +134,22 @@ class ProfileManager(object):
         self.peers = []
 
     def load_profile(self, profile_path):
-        """ Load profile from a json file
+        """ Load profile from a file
+
+        Open the pickle file, deserialize the content and
+        load it back into the profile manager.
         """
         avalon.dbgInfo('Loading profile from: {}'.format(profile_path))
-        with open(profile_path, 'r') as wgc_config:
-            profile = json.load(wgc_config)
-            wgc_config.close()
-
-        for peer in profile['peers']:
-            address = profile['peers'][peer]['address']
-            public_address = profile['peers'][peer]['public_address']
-            listen_port = profile['peers'][peer]['listen_port']
-            private_key = profile['peers'][peer]['private_key']
-            keep_alive = profile['peers'][peer]['keep_alive']
-            pm.peers.append(Peer(address, public_address, listen_port, private_key, keep_alive))
+        with open(profile_path, 'rb') as profile:
+            pm.peers = pickle.load(profile)
+            profile.close()
 
     def save_profile(self, profile_path):
-        """ Save current profile to a json file
-        """
+        """ Save current profile to a file
 
-        # Convert peer objects into dictionary format
-        profile = {}
-        profile['peers'] = {}
-        for peer in pm.peers:
-            profile['peers'][peer.address] = {}
-            profile['peers'][peer.address]['address'] = peer.address
-            profile['peers'][peer.address]['public_address'] = peer.public_address
-            profile['peers'][peer.address]['listen_port'] = peer.listen_port
-            profile['peers'][peer.address]['private_key'] = peer.private_key
-            profile['peers'][peer.address]['keep_alive'] = peer.keep_alive
+        Serializes the current profile with pickle
+        and dumps it into a file.
+        """
 
         # If profile already exists (file or link), ask the user if
         # we should overwrite it.
@@ -171,9 +166,9 @@ class ProfileManager(object):
 
         # Finally, write the profile into the destination file
         avalon.dbgInfo('Writing profile to: {}'.format(profile_path))
-        with open(profile_path, 'w') as wgc_config:
-            json.dump(profile, wgc_config, indent=2)
-            wgc_config.close()
+        with open(profile_path, 'wb') as profile:
+            pickle.dump(pm.peers, profile)
+            profile.close()
 
     def new_profile(self):
         """ Create new profile and flush the peers list