From 9a22321799a7dfe53ea679522582e6722c8264bb Mon Sep 17 00:00:00 2001
From: dadav <33197631+dadav@users.noreply.github.com>
Date: Wed, 1 Apr 2020 08:17:22 +0200
Subject: [PATCH] save in dotted format

---
 bin/pwnagotchi                       |  4 ++-
 pwnagotchi/__init__.py               |  1 +
 pwnagotchi/plugins/__init__.py       | 11 ++++++++-
 pwnagotchi/plugins/default/webcfg.py |  9 +++----
 pwnagotchi/utils.py                  | 37 ++++++++++++++++++++++++++++
 5 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/bin/pwnagotchi b/bin/pwnagotchi
index 5b93d3f..0b547ec 100755
--- a/bin/pwnagotchi
+++ b/bin/pwnagotchi
@@ -17,6 +17,7 @@ from pwnagotchi.agent import Agent
 from pwnagotchi.ui.display import Display
 from pwnagotchi import restart
 from pwnagotchi import fs
+from pwnagotchi.utils import DottedTomlEncoder
 
 
 def do_clear(display):
@@ -122,9 +123,10 @@ if __name__ == '__main__':
 
     config = utils.load_config(args)
     if args.print_config:
-        print(toml.dumps(config))
+        print(toml.dumps(config, encoder=DottedTomlEncoder()))
         sys.exit(0)
 
+    pwnagotchi.config = config
     fs.setup_mounts(config)
     log.setup_logging(args, config)
 
diff --git a/pwnagotchi/__init__.py b/pwnagotchi/__init__.py
index 185312a..58e7702 100644
--- a/pwnagotchi/__init__.py
+++ b/pwnagotchi/__init__.py
@@ -10,6 +10,7 @@ from pwnagotchi import fs
 from pwnagotchi._version import __version__
 
 _name = None
+config = None
 
 
 def set_name(new_name):
diff --git a/pwnagotchi/plugins/__init__.py b/pwnagotchi/plugins/__init__.py
index c8a70c6..5c2a662 100644
--- a/pwnagotchi/plugins/__init__.py
+++ b/pwnagotchi/plugins/__init__.py
@@ -5,6 +5,9 @@ import threading
 import importlib, importlib.util
 import logging
 from pwnagotchi.ui import view
+from pwnagotchi import config
+from pwnagotchi.utils import save_config
+
 
 default_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "default")
 loaded = {}
@@ -36,11 +39,17 @@ def toggle_plugin(name, enable=True):
 
     returns True if changed, otherwise False
     """
-    global loaded, database
+    global loaded, database, config
+
+    if config:
+        config['main']['plugins'][name] = enable
+        save_config(config, '/etc/pwnagotchi/config.toml')
+
     if not enable and name in loaded:
         if getattr(loaded[name], 'on_unload', None):
             loaded[name].on_unload(view.ROOT)
         del loaded[name]
+
         return True
 
     if enable and name in database and name not in loaded:
diff --git a/pwnagotchi/plugins/default/webcfg.py b/pwnagotchi/plugins/default/webcfg.py
index d870953..ba79078 100644
--- a/pwnagotchi/plugins/default/webcfg.py
+++ b/pwnagotchi/plugins/default/webcfg.py
@@ -2,8 +2,8 @@ import logging
 import json
 import toml
 import _thread
-import pwnagotchi.plugins as plugins
-from pwnagotchi import restart
+from pwnagotchi import restart, plugins
+from pwnagotchi.utils import save_config
 from flask import abort
 from flask import render_template_string
 
@@ -500,9 +500,8 @@ class WebConfig(plugins.Plugin):
             if path == "save-config":
                 try:
                     parsed_toml = toml.loads(request.get_json())
-                    with open('/etc/pwnagotchi/config.toml') as config_file:
-                        toml.dump(parsed_toml, config_file)
-
+                    save_config(parsed_toml, '/etc/pwnagotchi/config.toml')
+                    
                     _thread.start_new_thread(restart, (self.mode,))
                     return "success"
                 except Exception as ex:
diff --git a/pwnagotchi/utils.py b/pwnagotchi/utils.py
index 17bd79f..92a42e0 100644
--- a/pwnagotchi/utils.py
+++ b/pwnagotchi/utils.py
@@ -10,11 +10,43 @@ import json
 import shutil
 import toml
 import sys
+import re
 
 import pwnagotchi
+from toml.encoder import TomlEncoder, _dump_str
 from pwnagotchi.fs import ensure_write
 
 
+class DottedTomlEncoder(TomlEncoder):
+    """
+    Dumps the toml into the dotted-key format
+    """
+
+    def __init__(self, _dict=dict):
+        super(DottedTomlEncoder, self).__init__(_dict)
+
+    def dump_sections(self, o, sup):
+        retstr = ""
+        pre = ""
+
+        if sup:
+            pre = sup + "."
+
+        for section, value in o.items():
+            section = str(section)
+            qsection = section
+            if not re.match(r'^[A-Za-z0-9_-]+$', section):
+                qsection = _dump_str(section)
+            if value is not None:
+                if isinstance(value, dict):
+                    toadd, _ = self.dump_sections(value, pre + qsection)
+                    retstr += toadd
+                else:
+                    retstr += (pre + qsection + " = " +
+                                str(self.dump_value(value)) + '\n')
+        return (retstr, self._dict())
+
+
 # https://stackoverflow.com/questions/823196/yaml-merge-in-python
 def merge_config(user, default):
     if isinstance(user, dict) and isinstance(default, dict):
@@ -44,6 +76,11 @@ def keys_to_str(data):
 
     return converted_dict
 
+def save_config(config, target):
+    with open(target, 'wt') as fp:
+        fp.write(toml.dumps(config, encoder=DottedTomlEncoder()))
+    return True
+
 def load_config(args):
     default_config_path = os.path.dirname(args.config)
     if not os.path.exists(default_config_path):