Add proxy based on env and fix iterfzf error

you-n-g · Aug 11, 2020 · 78ff2e0 · 78ff2e0
1 parent 0eb6f36
commit 78ff2e0
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 3 deletions.
diff --git a/wan/__init__.py b/wan/__init__.py
@@ -8,10 +8,11 @@
 import psutil
 import time
 import sys
+import os
 
 
 class Notifier:
-    def __init__(self, config_path: str = '~/.dotfiles/.notifers.yaml'):
+    def __init__(self, config_path: str = '~/.dotfiles/.notifiers.yaml'):
         """__init__.
 
         Parameters
@@ -25,6 +26,15 @@ def __init__(self, config_path: str = '~/.dotfiles/.notifers.yaml'):
                 chat_id: <Your Chat ID>
                 token: <Your token>
             ```
+
+            If you need proxy for your provider, please use the config below.
+            The env will be updated when running `ntf` method
+            [This solution](https://github.com/liiight/notifiers/issues/236) is proposed by notifiers
+            ```
+            env:
+                HTTP_PROXY: 'http://IP:PORT'
+                HTTPS_PROXY: 'http://IP:PORT'
+            ```
         """
         # TODO: DEBUG mode
         path = Path(config_path).expanduser()
@@ -43,10 +53,23 @@ def __init__(self, config_path: str = '~/.dotfiles/.notifers.yaml'):
         kwargs = self.config['kwargs']
         self._ntf = partial(self._provider.notify, **kwargs)
 
+        self.env = self.config.get('env', {})
+
     def ntf(self, *messages):
         message = " ".join(messages)
         logger.debug("Sending message: {}".format(message))
+        if len(message) == 0:
+            logger.warning("Blank message.")
+
+        # set proxy if needed
+        env_back = os.environ.copy()
+        os.environ.update(self.env)
         self._ntf(message=message)
+        for k, v in self.env.items():
+            if k not in env_back:
+                del os.environ[k]
+            else:
+                os.environ[k] = env_back[k]
 
     @staticmethod
     def _get_process_info(pid):

diff --git a/wan/utils.py b/wan/utils.py
@@ -1,6 +1,8 @@
 import subprocess
 from iterfzf import iterfzf
 from loguru import logger
+import os
+import stat
 
 
 def iter_ps():
@@ -16,10 +18,16 @@ def get_pid_from_line(line):
 
 def get_pid_via_fzf(exact=True):
     try:
+        # FIXME: This is a bug from iterfzf. The fzf may not be executable
         selected_line = iterfzf(iter_ps(), multi=False, exact=exact)
     except PermissionError as e:
-        logger.error(f'Please make {e.filename} executable(e.g  `chmod a+x {e.filename}`).')
-        return None
+        try:
+            os.chmod(e.filename, os.stat(e.filename).st_mode |  stat.S_IEXEC)
+        except PermissionError:
+            logger.error(f'Please make {e.filename} executable(e.g  `chmod a+x {e.filename}`).')
+            return None
+        else:
+            selected_line = iterfzf(iter_ps(), multi=False, exact=exact)
     return get_pid_from_line(selected_line)