Merge pull request #98 from NamamiShanker/notion

🚀 Add Notion.so login feature to get token

Author: GouravSardana

README.md

@@ -43,6 +43,7 @@ Options: <br>
 `-st, --start -> Introduces Dynamic CLI` <br>
 `-v, --version -> Gives the Version of the CLI` <br>
 `-s, --search -> Search a question on Stackoverflow` <br>
+`-no, --notion -> Open browser to login to Notion.so` <br>
 `-d, --debug -> Turn on Debugging mode` <br>
 `-c, --custom -> Setup a custom API key` <br>
 `-h, --help -> Shows this message and exit` <br>

main.py

@@ -55,6 +55,14 @@
                     help="View and organise the playbook",
                     action='store_true')
 
+parser.add_argument("-no",
+                    "--notion",
+                    help="\
+                    Login to your Notion account to save playbook.\
+                    Opens a browser window for you to login to\
+                    your Notion accout",
+                    action='store_true')
+
 ARGV = parser.parse_args()
 
 search_flag = Search(ARGV)

src/arguments/error.py

@@ -18,3 +18,20 @@ def evoke_search_error(self, error_statement):
         ]
         for text_to_print in print_text:
             print(text_to_print)
+
+class LoginError():
+    def __init__(self, error_statement, success=False):
+        """
+        Implements error printing for User Login
+        :error_statement: Error statement to print
+        :success: Indicates success of login attempt
+                  Prints in green if True else red
+        """
+        self.error_statement = error_statement
+        self.success = success
+        self.evoke_search_error()
+
+    def evoke_search_error(self):
+        color = "green" if self.success else "red"
+        print_text = colored(self.error_statement, color)
+        print(print_text)

src/arguments/notion.py

@@ -0,0 +1,87 @@
+import os
+
+from .utility import get_browser_driver
+from .error import LoginError
+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
+
+def get_token_from_cookie(cookie, token):
+    for el in cookie:
+        if el['name'] == token:
+            return el
+
+def get_token_from_file():
+    try:
+        with open(TOKEN_FILE_PATH, 'r') as f:
+            data = f.read()
+    except Exception as e:
+        print(e)
+    if(not data or data==""):
+        raise RuntimeError("Token not found in file")
+    else:
+        return data
+
+def get_cookies_from_login():
+    """Capture browser cookies for authentication."""
+    driver = get_browser_driver()
+    try:
+        driver.get(LOGIN_PATH)
+        WebDriverWait(driver, 300).until(
+            EC.presence_of_element_located((By.CLASS_NAME,
+                                            "notion-sidebar-container")))
+        return driver.get_cookies()
+    except Exception as e:
+        print(e)
+    finally:
+        driver.quit()
+
+class NotionClient():
+
+    """
+    Implements Login and token retrieval.
+
+    Handles the entire procedure of connecting to User's Notion account,
+    generating Notion's tokenv2_cookie, storing it locally and uploading
+    content to User's space
+    """
+    def __init__(self):
+        """
+        No input parameters required for instatiating object.
+
+        :tokenv2_cookie: stores the cookie containing user's tokenv2
+        :tokenv2_key: used to create environment variable
+        """
+        self.tokenv2_cookie = None
+        self.tokenv2_key = 'TOKENV2'
+
+    def save_token_file(self):
+        if(self.tokenv2_cookie):
+            with open(TOKEN_FILE_PATH, 'w') as f:
+                f.write(str(self.tokenv2_cookie))
+
+    def get_tokenv2_cookie(self):
+        message_success = "Successfully logged into Notion \U0001F389"
+        message_failure = "Login unsuccessful. Please try again \U0001F615"
+        # Sets 'tokenv2_cookie equal to the particular cookie containing token_v2
+        if not self.tokenv2_cookie:
+            try:
+                self.tokenv2_cookie = get_token_from_file()
+                LoginError(message_success, success=True)
+            except Exception:
+                try:
+                    cookies = get_cookies_from_login()
+                    self.tokenv2_cookie = get_token_from_cookie(cookies, 'token_v2')
+                except Exception:
+                    self.tokenv2_cookie = None
+                finally:
+                    if self.tokenv2_cookie:
+                        os.environ[self.tokenv2_key] = str(self.tokenv2_cookie)
+                        LoginError(message_success, success=True)
+                        self.save_token_file()
+                    else:
+                        LoginError(message_failure, success=False)
+        return self.get_tokenv2_cookie

src/arguments/search.py

@@ -9,6 +9,7 @@
 from .save import SaveSearchResults
 from .update import UpdateApplication
 from .api_test import ApiTesting
+from .notion import NotionClient
 
 version = "0.1.0"
 class Prompt():
@@ -53,6 +54,8 @@ def search_args(self):
             update.check_for_updates()
         elif self.arguments.GET:
             self.api_test_object.get_request()
+        elif self.arguments.notion:
+            NotionClient().get_tokenv2_cookie()
 
     def search_for_results(self, save=False):
         queries = ["What do you want to search", "Tags"]

src/arguments/settings.py

@@ -0,0 +1,18 @@
+import os
+from pathlib import Path
+
+NOTION_URL = "https://www.notion.so/"
+LOGIN_PATH = NOTION_URL + "/login"
+DATA_DIR = os.environ.get(
+    "DYNAMIC_DATA_DIR", str(Path(os.path.expanduser("~")).joinpath(".dynamic"))
+)
+PLAYBOOK_FILE = str(Path(DATA_DIR).joinpath("playbook.json"))
+TOKEN_FILE_PATH = str(Path(DATA_DIR).joinpath("tokenv2_cookie"))
+
+try:
+    os.makedirs(DATA_DIR)
+except FileExistsError:
+    pass
+
+if(not os.path.isfile(TOKEN_FILE_PATH)):
+    open(TOKEN_FILE_PATH, 'w').close()

src/arguments/utility.py

@@ -15,6 +15,7 @@
 from .error import SearchError
 from .save import SaveSearchResults
 from .markdown import MarkdownRenderer
+from .settings import PLAYBOOK_FILE
 
 # Required for OAuth
 import json
@@ -34,25 +35,32 @@
 
 console = Console()
 
+def get_browser_driver():
+    # 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")
+
 class Playbook():
     def __init__(self):
-        self.linux_path = "/home/{}/Documents/dynamic".\
-                          format(os.getenv('USER'))
-        self.mac_path = "/Users/{}/Documents/dynamic".\
-                        format(os.getenv('USER'))
-        self.file_name = 'dynamic_playbook.json'
         self.key = 'DYNAMIC'
 
     @property
     def playbook_path(self):
         """Create an environment variable 'DYNAMIC containing the path of dynamic_playbook.json and returns i."""
         if not os.getenv(self.key):
-            if(sys.platform=='linux'):
-                os.environ[self.key] = os.path.\
-                                       join(self.linux_path, self.file_name)
-            elif(sys.platform=='darwin'):
-                os.environ[self.key] = os.path.\
-                                       join(self.mac_path, self.file_name)
+            os.environ[self.key] = PLAYBOOK_FILE
         return os.getenv(self.key)
 
     @property
@@ -380,21 +388,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 = get_browser_driver()
 
         # Open auth_url in one of the supported browsers
         driver.get(auth_url)