From f6605640b30585fcca7e8e5579cc8f1db8220c14 Mon Sep 17 00:00:00 2001
From: Simone Margaritelli <evilsocket@gmail.com>
Date: Thu, 3 Oct 2019 14:45:39 +0200
Subject: [PATCH] custom config override file

---
 sdcard/rootfs/root/pwnagotchi/config.yml      |  2 +-
 sdcard/rootfs/root/pwnagotchi/scripts/main.py | 78 ++++++++++---------
 .../scripts/pwnagotchi/mesh/utils.py          |  6 +-
 .../pwnagotchi/scripts/pwnagotchi/utils.py    | 24 ++++++
 4 files changed, 70 insertions(+), 40 deletions(-)
 create mode 100644 sdcard/rootfs/root/pwnagotchi/scripts/pwnagotchi/utils.py

diff --git a/sdcard/rootfs/root/pwnagotchi/config.yml b/sdcard/rootfs/root/pwnagotchi/config.yml
index cc2ea20..b71ddde 100644
--- a/sdcard/rootfs/root/pwnagotchi/config.yml
+++ b/sdcard/rootfs/root/pwnagotchi/config.yml
@@ -95,7 +95,7 @@ ui:
     # IMPORTANT: The lifespan of an eINK display depends on the cumulative amount of refreshes. If you want to
     # preserve your display over time, you should set this value to 0.0 so that the display will be refreshed only
     # if any of the important data fields changed (the uptime and blinking cursor won't trigger a refresh).
-    fps: 0.3
+    fps: 0.0
     display:
         enabled: true
         rotation: 180
diff --git a/sdcard/rootfs/root/pwnagotchi/scripts/main.py b/sdcard/rootfs/root/pwnagotchi/scripts/main.py
index 09f6ec3..3bf19c7 100755
--- a/sdcard/rootfs/root/pwnagotchi/scripts/main.py
+++ b/sdcard/rootfs/root/pwnagotchi/scripts/main.py
@@ -1,11 +1,16 @@
 #!/usr/bin/python3
+import os
 import argparse
-import yaml
 import time
 import traceback
 
+import yaml
+
 import core
-import pwnagotchi, pwnagotchi.plugins as plugins
+import pwnagotchi
+import pwnagotchi.utils as utils
+import pwnagotchi.version as version
+import pwnagotchi.plugins as plugins
 
 from pwnagotchi.log import SessionParser
 from pwnagotchi.voice import Voice
@@ -14,53 +19,52 @@ from pwnagotchi.ui.display import Display
 
 parser = argparse.ArgumentParser()
 
-parser.add_argument('-C', '--config', action='store', dest='config', default='/root/pwnagotchi/config.yml')
+parser.add_argument('-C', '--config', action='store', dest='config', default='/root/pwnagotchi/config.yml',
+                    help='Main configuration file.')
+parser.add_argument('-U', '--user-config', action='store', dest='user_config', default='/root/pwnagotchi.yml',
+                    help='If this file exists, configuration will be merged and this will override default values.')
 
 parser.add_argument('--manual', dest="do_manual", action="store_true", default=False, help="Manual mode.")
 parser.add_argument('--clear', dest="do_clear", action="store_true", default=False,
                     help="Clear the ePaper display and exit.")
 
 args = parser.parse_args()
+config = utils.load_config(args)
 
 if args.do_clear:
     print("clearing the display ...")
-    with open(args.config, 'rt') as fp:
-        config = yaml.safe_load(fp)
-        cleardisplay = config['ui']['display']['type']
-        if cleardisplay in ('inkyphat', 'inky'):
-            print("inky display")
-            from inky import InkyPHAT
+    cleardisplay = config['ui']['display']['type']
+    if cleardisplay in ('inkyphat', 'inky'):
+        print("inky display")
+        from inky import InkyPHAT
 
-            epd = InkyPHAT(config['ui']['display']['color'])
-            epd.set_border(InkyPHAT.BLACK)
-            self._render_cb = self._inky_render
-        elif cleardisplay in ('papirus', 'papi'):
-            print("papirus display")
-            from pwnagotchi.ui.papirus.epd import EPD
+        epd = InkyPHAT(config['ui']['display']['color'])
+        epd.set_border(InkyPHAT.BLACK)
+        self._render_cb = self._inky_render
+    elif cleardisplay in ('papirus', 'papi'):
+        print("papirus display")
+        from pwnagotchi.ui.papirus.epd import EPD
 
-            os.environ['EPD_SIZE'] = '2.0'
-            epd = EPD()
-            epd.clear()
-        elif cleardisplay in ('waveshare_1', 'ws_1', 'waveshare1', 'ws1'):
-            print("waveshare v1 display")
-            from pwnagotchi.ui.waveshare.v1.epd2in13 import EPD
+        os.environ['EPD_SIZE'] = '2.0'
+        epd = EPD()
+        epd.clear()
+    elif cleardisplay in ('waveshare_1', 'ws_1', 'waveshare1', 'ws1'):
+        print("waveshare v1 display")
+        from pwnagotchi.ui.waveshare.v1.epd2in13 import EPD
 
-            epd = EPD()
-            epd.init(epd.lut_full_update)
-            epd.Clear(0xFF)
-        elif cleardisplay in ('waveshare_2', 'ws_2', 'waveshare2', 'ws2'):
-            print("waveshare v2 display")
-            from pwnagotchi.ui.waveshare.v2.waveshare import EPD
+        epd = EPD()
+        epd.init(epd.lut_full_update)
+        epd.Clear(0xFF)
+    elif cleardisplay in ('waveshare_2', 'ws_2', 'waveshare2', 'ws2'):
+        print("waveshare v2 display")
+        from pwnagotchi.ui.waveshare.v2.waveshare import EPD
 
-            epd = EPD()
-            epd.init(epd.FULL_UPDATE)
-            epd.Clear(0xff)
-        else:
-            print("unknown display type %s" % cleardisplay)
-        quit()
-
-with open(args.config, 'rt') as fp:
-    config = yaml.safe_load(fp)
+        epd = EPD()
+        epd.init(epd.FULL_UPDATE)
+        epd.Clear(0xff)
+    else:
+        print("unknown display type %s" % cleardisplay)
+    quit()
 
 plugins.load_from_path(plugins.default_path)
 if 'plugins' in config['main'] and config['main']['plugins'] is not None:
@@ -71,7 +75,7 @@ plugins.on('loaded')
 display = Display(config=config, state={'name': '%s>' % pwnagotchi.name()})
 agent = Agent(view=display, config=config)
 
-core.log("%s@%s (v%s)" % (pwnagotchi.name(), agent._identity, pwnagotchi.version))
+core.log("%s@%s (v%s)" % (pwnagotchi.name(), agent._identity, version.version))
 # for key, value in config['personality'].items():
 #    core.log("  %s: %s" % (key, value))
 
diff --git a/sdcard/rootfs/root/pwnagotchi/scripts/pwnagotchi/mesh/utils.py b/sdcard/rootfs/root/pwnagotchi/scripts/pwnagotchi/mesh/utils.py
index 8c84b94..92e4939 100644
--- a/sdcard/rootfs/root/pwnagotchi/scripts/pwnagotchi/mesh/utils.py
+++ b/sdcard/rootfs/root/pwnagotchi/scripts/pwnagotchi/mesh/utils.py
@@ -1,7 +1,9 @@
 import _thread
 
 import core
-import pwnagotchi, pwnagotchi.plugins as plugins
+import pwnagotchi
+import pwnagotchi.version as version
+import pwnagotchi.plugins as plugins
 from pwnagotchi.mesh import get_identity
 
 
@@ -22,7 +24,7 @@ class AsyncAdvertiser(object):
         self._advertiser = Advertiser(
             self._config['main']['iface'],
             pwnagotchi.name(),
-            pwnagotchi.version,
+            version.version,
             self._identity,
             period=0.3,
             data=self._config['personality'])
diff --git a/sdcard/rootfs/root/pwnagotchi/scripts/pwnagotchi/utils.py b/sdcard/rootfs/root/pwnagotchi/scripts/pwnagotchi/utils.py
new file mode 100644
index 0000000..358b2fe
--- /dev/null
+++ b/sdcard/rootfs/root/pwnagotchi/scripts/pwnagotchi/utils.py
@@ -0,0 +1,24 @@
+import yaml
+import os
+
+# https://stackoverflow.com/questions/823196/yaml-merge-in-python
+def merge_config(user, default):
+    if isinstance(user, dict) and isinstance(default, dict):
+        for k, v in default.items():
+            if k not in user:
+                user[k] = v
+            else:
+                user[k] = merge_config(user[k], v)
+    return user
+
+
+def load_config(args):
+    with open(args.config, 'rt') as fp:
+        config = yaml.safe_load(fp)
+
+    if os.path.exists(args.user_config):
+        with open(args.user_config, 'rt') as fp:
+            user_config = yaml.safe_load(fp)
+            config = merge_config(user_config, config)
+
+    return config
\ No newline at end of file