Add interfaces group + button to quick up all interfaces

README.md

@@ -8,13 +8,20 @@ wg-tray enables to quickly bring up/down wireguard interfaces from the system tr
 
 ## Usage
 ```bash
-wg-tray [-h] [-v] [-c CONFIG]
+usage: wg-tray [-h] [-v] [-c CONFIG] [-g CONFIG_GROUPS]
 
-optional arguments:
-  -h, --help                   show this help message and exit
-  -v, --version                show program version info and exit
-  -c CONFIG, --config CONFIG   path to the config file listing all wireguard interfaces 
-                               (default: none; use root privileges to look up in /etc/wireguard/)
+A simple UI tool to handle WireGuard interfaces
+
+options:
+  -h, --help            show this help message and exit
+  -v, --version         show program version info and exit
+  -c CONFIG, --config CONFIG
+                        path to the config file listing all WireGuard interfaces (if none is
+                        provided, use root privileges to look up in /etc/wireguard/) (default:
+                        None)
+  -g CONFIG_GROUPS, --config-groups CONFIG_GROUPS
+                        Path to the config (.ini file) to have groups of wireguard configs.
+                        (default: ~/.wireguard/wg_tray_groups.ini)
 ```
 The config file should simply list all wireguard interfaces, separated either by newlines or spaces (e.g. `wg0 wg1` or
 ```
@@ -29,3 +36,29 @@ If you want to avoid being prompted for your root password each time you run `wg
 <username> ALL=(ALL) NOPASSWD: /usr/bin/wg
 <username> ALL=(ALL) NOPASSWD: /usr/bin/wg-quick
 ```
+
+### Config groups
+You can define configs groups in the config groups files (located `./wireguard/wg_tray_groups.ini` by default).
+
+#### Example
+Here's an example of a `wg_tray_groups.ini` config file:
+
+```ini
+[settings]
+pick_one_at_random = false
+
+[Group 01]
+pick_one_at_random = true
+interfaces = inter01
+           = inter02
+           = inter03
+           = inter04
+
+[Group 02]
+interfaces = inter05
+           = inter06
+```
+
+#### Configs groups settings: `pick_one_at_random`
+If you whish to only up one random interfaces from the group, you can define the settings `pick_one_at_random` in your section.
+You can also use it on all groups be defining a `settings` section.

wgtray/actions/interface.py

@@ -1,4 +1,7 @@
+from itertools import chain
+import logging
 import pathlib
+import random
 import subprocess
 import threading
 
@@ -10,6 +13,8 @@
 
 RES_PATH = pathlib.Path(__file__).parent.parent.resolve() / "res"
 
+logger = logging.getLogger(__name__)
+
 
 class WGInterface(QAction):
     done = pyqtSignal(bool, str, name="done")  # necessary to put outside of __init__
@@ -64,3 +69,93 @@ def bring_up_down(self, is_up):
             err_msg = std_err.decode()
 
         self.done.emit(stat, err_msg)
+
+
+class WGInterfaceAll(QAction):
+    done = pyqtSignal(int, name="done")  # necessary to put outside of __init__
+
+    def __init__(self, name, parent, interfaces, type_, subgroups=None, pick_one_at_random=False, refresh=None):
+        super().__init__(name, parent)
+
+        self.name = name
+        self.parent = parent
+        self.interfaces = interfaces
+        self.type_ = type_
+        self.refresh = refresh
+
+        self.subgroups = subgroups if subgroups else []
+        self.pick_one_at_random = pick_one_at_random
+
+        self.triggered.connect(self.toggle)
+        self.done.connect(self._done)
+        self.updateIcon()
+
+    def updateIcon(self):
+        if self.type_:
+            icon_path = f"{RES_PATH}/green_arrow_up.png"
+        else:
+            icon_path = f"{RES_PATH}/grey_arrow_down.png"
+        self.setIcon(QIcon(icon_path))
+
+    @pyqtSlot(int)
+    def _done(self, count):
+
+        pt = "upped" if self.type_ else "downed"
+        if count:
+            self.parent.tray.showMessage("Informations", f"Successfully {pt} {count} interface(s)", QSystemTrayIcon.NoIcon)
+
+        else:
+            self.parent.tray.showMessage("Warning", f"No interfaces where {pt}", QSystemTrayIcon.NoIcon)
+
+        self.loadingSpinner.stop()
+        self.updateIcon()
+
+        if self.refresh:
+            self.refresh()
+
+    def toggle(self):
+        # Loading animation
+        self.loadingSpinner = QMovie(f"{RES_PATH}/loader.gif")
+        self.loadingSpinner.frameChanged.connect(lambda: self.setIcon(QIcon(self.loadingSpinner.currentPixmap())))
+        self.loadingSpinner.start()
+
+        # Launch command
+        t = threading.Thread(target=self._toggle)
+        t.start()
+
+    def get_iterfaces_to_workon(self):
+        """Return the list of interfaces to down/up."""
+
+        if self.subgroups:
+            return chain.from_iterable(group.get_iterfaces_to_workon() for group in self.subgroups)
+
+        if self.pick_one_at_random:
+            return random.choices(self.interfaces, k=1)
+
+        return self.interfaces
+
+    def _toggle(self):
+        kw = "up" if self.type_ else "down"
+
+        def _interface_open(interface):
+            subp = subprocess.Popen(f"sudo wg-quick {kw} {interface}", stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
+
+            _, std_err = subp.communicate()  # blocking
+            if subp.returncode == 0:
+                return True, ""
+
+            return False, std_err.decode()
+
+        upped_interface = 0
+
+        for interface in self.get_iterfaces_to_workon():
+            success, err_msg = _interface_open(interface)
+
+            if success:
+                logger.info(f"Interface: {interface}, sucessfully mounted")
+                upped_interface += 1
+
+            else:
+                logger.info(f"Interface: {interface}, error while mounting: {err_msg}")
+
+        self.done.emit(upped_interface)

wgtray/wgtray.py

@@ -1,27 +1,29 @@
+from configparser import ConfigParser
 import argparse
+import logging
 import os
 import pathlib
 import signal
 import sys
 
 
 from PyQt5.QtCore import pyqtSlot, QCoreApplication, QTimer
-from PyQt5.QtGui import QContextMenuEvent, QIcon, QMovie
+from PyQt5.QtGui import QIcon, QMovie
 from PyQt5.QtWidgets import QAction, QApplication, QMenu, QSystemTrayIcon
 
 
 from . import __description__, __version__
-from .actions.interface import WGInterface
+from .actions.interface import WGInterface, WGInterfaceAll
 
 
 RES_PATH = pathlib.Path(__file__).parent.resolve() / "res"
 
 
 class WGTrayIcon(QSystemTrayIcon):
-    def __init__(self, config_path=None):
+    def __init__(self, config_path=None, config_menu=None):
         super().__init__()
 
-        self.menu = WGMenu(self, config_path)
+        self.menu = WGMenu(self, config_path, config_menu)
         self.setContextMenu(self.menu)
         self.activated.connect(self.activateMenu)  # show also on left click
 
@@ -36,22 +38,123 @@ def activateMenu(self, activationReason):
 
 
 class WGMenu(QMenu):
-    def __init__(self, tray, config_path):
+    def __init__(self, tray, config_path, config_menu=None):
         super().__init__()
 
         self.tray = tray
+        self.config_menu = config_menu
+
+        itfs_up = self.read_status()
+
+        interfaces_done = []
+        interface_actions = []
+        self.menus = []
+
+        if self.config_menu:
+            general_pick_one_at_random = False
+            if "settings" in config:
+                settings = self.config_menu["settings"]
+                general_pick_one_at_random = settings.get("pick_one_at_random", "false") == "true"
+
+            for section in self.config_menu.sections():
+                if section == "settings":
+                    continue
+
+                if "pick_one_at_random" in self.config_menu[section]:
+                    pick_one_at_random = self.config_menu[section].get("pick_one_at_random", "false") == "true"
+
+                else:
+                    pick_one_at_random = general_pick_one_at_random
+
+                menu = self.addMenu(str(section))
+                self.menus.append(menu)
+
+                if not menu:
+                    continue
+
+                section_interfaces = []
+                for interface in self.config_menu[section]["interfaces"].strip().split():
+                    action = WGInterface(interface, self, interface in itfs_up)
+                    action.updateIcon()
+                    menu.addAction(action)
+
+                    section_interfaces.append(interface)
+
+                interfaces_done.extend(section_interfaces)
+
+                menu.addSeparator()
+                interface_action = WGInterfaceAll(
+                    "Up one random interface" if pick_one_at_random else "Up all interfaces",
+                    self,
+                    section_interfaces,
+                    True,
+                    pick_one_at_random=pick_one_at_random,
+                    refresh=self.startRefresh,
+                )
+                menu.addAction(interface_action)
+                interface_actions.append(interface_action)
+
+                menu.addAction(
+                    WGInterfaceAll(
+                        "Down all interfaces",
+                        self,
+                        section_interfaces,
+                        False,
+                        # We want to down all interfaces, regarding of the settings for the up
+                        pick_one_at_random=False,
+                        refresh=self.startRefresh,
+                    )
+                )
 
         if config_path:
             with open(config_path) as f:
                 itfs = f.read()
         else:
             itfs = os.popen("sudo ls /etc/wireguard | grep .conf | awk -F \".\" '{print $1}'").read()
-        itfs_up = self.read_status()
 
-        for itf_name in itfs.strip().split():
-            action = WGInterface(itf_name, self, itf_name in itfs_up)
-            action.updateIcon()
-            self.addAction(action)
+        itfs = itfs.strip().split()
+
+        for itf_name in itfs:
+            if itf_name not in interfaces_done:
+                action = WGInterface(itf_name, self, itf_name in itfs_up)
+                action.updateIcon()
+                self.addAction(action)
+
+                interfaces_done.append(itf_name)
+
+        self.addSeparator()
+
+        self.addAction(
+            WGInterfaceAll(
+                "Up interfaces on all groups",
+                self,
+                [],
+                True,
+                subgroups=interface_actions,
+                pick_one_at_random=False,
+                refresh=self.startRefresh,
+            )
+        )
+        self.addAction(
+            WGInterfaceAll(
+                "Up all interfaces",
+                self,
+                interfaces_done,
+                True,
+                pick_one_at_random=False,
+                refresh=self.startRefresh,
+            )
+        )
+        self.addAction(
+            WGInterfaceAll(
+                "Down interfaces on all groups",
+                self,
+                interfaces_done,
+                False,
+                pick_one_at_random=False,
+                refresh=self.startRefresh,
+            )
+        )
 
         self.addSeparator()
 
@@ -84,6 +187,12 @@ def reloadStatus(self):
                 action.setUp(action.text() in itfs_up)
                 action.updateIcon()
 
+        for menu in self.menus:
+            for action in menu.actions():
+                if isinstance(action, WGInterface):
+                    action.setUp(action.text() in itfs_up)
+                    action.updateIcon()
+
     def showTearOff(self):
         self.showTearOffMenu(QApplication.desktop().availableGeometry().center())
         self.closeTearOff.setVisible(True)
@@ -142,7 +251,10 @@ def quit(self):
 
 def parse_args():
 
-    parser = argparse.ArgumentParser(description=__description__)
+    parser = argparse.ArgumentParser(
+        description=__description__,
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
     parser.add_argument("-v", "--version", help="show program version info and exit", action="version", version=__version__)
     parser.add_argument(
         "-c",
@@ -151,17 +263,29 @@ def parse_args():
         default=None,
         type=str,
     )
+    parser.add_argument(
+        "-g",
+        "--config-groups",
+        help="Path to the config (.ini file) to have groups of wireguard configs.",
+        default="~/.wireguard/wg_tray_groups.ini",
+        type=str,
+    )
 
     args = parser.parse_args()
 
-    return args.config
+    return args
 
 
+logging.basicConfig(level=logging.INFO)
+config = ConfigParser()
+
 app = QApplication(sys.argv)
 
-config_path = parse_args()
+parser_args = parse_args()
+config.read(pathlib.Path(parser_args.config_groups).expanduser())
+
+WGTrayIcon(parser_args.config, config_menu=config)
 
-WGTrayIcon(config_path)
 
 signal.signal(signal.SIGINT, signal.SIG_DFL)