Add feature to get token_v2 from login

Author: NamamiShanker

src/arguments/notion.py

@@ -3,16 +3,13 @@
 
 from .utility import Utility
 from .error import SearchError
+from .settings import LOGIN_PATH
+from .settings import TOKEN_FILE_PATH
 
 from selenium.webdriver.common.by import By
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
 
-
-from webdriver_manager.chrome import ChromeDriverManager
-from webdriver_manager.firefox import GeckoDriverManager
-from webdriver_manager.microsoft import EdgeChromiumDriverManager
-
 def get_token_from_cookie(cookie, token):
     for el in cookie:
         if el['name'] == token:
@@ -27,21 +24,23 @@ class NotionClient():
     """
     def __init__(self):
         self.tokenv2_cookie = None
-        self.base_url = "https://www.notion.so/"
-        self.login_path = "login/"
-        self.linux_path = "/home/{}/Documents/dynamic".format(os.getenv('USER'))
-        self.mac_path = "/Users/{}/Documents/dynamic".format(os.getenv('USER'))
-        self.file_name = 'tokenv2_cookie.key'
+        self.tokenv2_key = 'TOKENV2'
 
     def get_token_from_file(self):
-        raise FileNotFoundError("File not found")
-        return None
-
-    def get_login_path(self):
-        return self.base_url + self.login_path
+        try:
+            with open(TOKEN_FILE_PATH, 'r') as f:
+                data = f.read()
+        except Exception as e:
+            raise e
+        if(not data or data==""):
+            raise Exception("Token not found in file")
+        else:
+            return data
 
     def save_token_file(self):
-        pass
+        if(self.tokenv2_cookie):
+            with open(TOKEN_FILE_PATH, 'w') as f:
+                f.write(str(self.tokenv2_cookie))
 
     def get_cookies_from_login(self):
         """
@@ -51,7 +50,7 @@ def get_cookies_from_login(self):
         """
         driver = Utility().get_browser_driver()
         try:
-            driver.get(self.get_login_path())
+            driver.get(LOGIN_PATH)
             WebDriverWait(driver, 300).until(
                 EC.presence_of_element_located((By.CLASS_NAME,
                                                 "notion-sidebar-container")))
@@ -61,7 +60,7 @@ def get_cookies_from_login(self):
         finally:
             driver.quit()
 
-    def set_tokenv2_cookie(self):
+    def get_tokenv2_cookie(self):
         # Sets 'tokenv2_cookie equal to the particular cookie containing token_v2
         if not self.tokenv2_cookie:
             try:
@@ -74,6 +73,9 @@ def set_tokenv2_cookie(self):
                     print(e)
                     self.tokenv2_cookie = None
                 finally:
-                    os.environ['tokenv2_cookie'] = self.tokenv2_cookie
                     if self.tokenv2_cookie:
-                        self.save_token_file()
+                        os.environ[self.tokenv2_key] = str(self.tokenv2_cookie)
+                        self.save_token_file()
+                        return self.get_tokenv2_cookie
+                    else:
+                        raise Exception("Cookie unreachable")

src/arguments/search.py

@@ -56,7 +56,7 @@ def search_args(self):
         elif self.arguments.GET:
             self.api_test_object.get_request()
         elif self.arguments.notion:
-            NotionClient().print_token_v2()
+            NotionClient().get_tokenv2_cookie()
 
     def search_for_results(self, save=False):
         queries = ["What do you want to search", "Tags"]

src/arguments/settings.py

@@ -7,9 +7,12 @@
     "DYNAMIC_DATA_DIR", str(Path(os.path.expanduser("~")).joinpath(".dynamic"))
 )
 PLAYBOOK_FILE = str(Path(DATA_DIR).joinpath("playbook.json"))
-TOKEN_COOKIE_FILE = str(Path(DATA_DIR).joinpath("tokenv2_cookie"))
+TOKEN_FILE_PATH = str(Path(DATA_DIR).joinpath("tokenv2_cookie"))
 
 try:
     os.makedirs(DATA_DIR)
 except FileExistsError:
-    pass
+    pass
+
+if(not os.path.isfile(TOKEN_FILE_PATH)):
+    open(TOKEN_FILE_PATH, 'w').close()

src/arguments/utility.py

@@ -15,7 +15,7 @@
 from .error import SearchError
 from .save import SaveSearchResults
 from .markdown import MarkdownRenderer
-from .settings import PLAYBOOK_DIR
+from .settings import PLAYBOOK_FILE
 
 # Required for OAuth
 import json
@@ -47,7 +47,7 @@ def playbook_path(self):
         """
         if not os.getenv(self.key):
             if(sys.platform=='linux'):
-                os.environ[self.key] = PLAYBOOK_DIR
+                os.environ[self.key] = PLAYBOOK_FILE
         return os.getenv(self.key)
 
     @property
@@ -351,6 +351,24 @@ def get_ans(self, questions_list):
         stackoverflow_panel.display_panel(questions_list)
         # Support for reddit searching can also be implemented from here
 
+    def get_browser_driver(self):
+        # Try to install web drivers for one of these browsers
+        # Chrome, Firefox, Edge (One of them must be installed)
+        try:
+            return webdriver.Chrome(ChromeDriverManager().install())
+        except ValueError:
+            try:
+                return webdriver.Firefox(executable_path=
+                                           GeckoDriverManager().install())
+            except ValueError:
+                try:
+                    return webdriver.Edge(EdgeChromiumDriverManager().\
+                                            install())
+                except ValueError:
+                    print("You do not have one of these supported browsers:" +
+                          "Chrome, Firefox, Edge")
+        pass
+
     # Get an access token and extract to a JSON file "access_token.json"
     @classmethod
     def setCustomKey(self):
@@ -375,21 +393,7 @@ def setCustomKey(self):
                                   scope=scopes, redirect_uri=redirect_uri)
         auth_url, state = stackApps.authorization_url(authorization_url)
 
-        # Try to install web drivers for one of these browsers
-        # Chrome, Firefox, Edge (One of them must be installed)
-        try:
-            driver = webdriver.Chrome(ChromeDriverManager().install())
-        except ValueError:
-            try:
-                driver = webdriver.Firefox(executable_path=
-                                           GeckoDriverManager().install())
-            except ValueError:
-                try:
-                    driver = webdriver.Edge(EdgeChromiumDriverManager().\
-                                            install())
-                except ValueError:
-                    print("You do not have one of these supported browsers:" +
-                          "Chrome, Firefox, Edge")
+        driver = self.get_browser_driver()
 
         # Open auth_url in one of the supported browsers
         driver.get(auth_url)