add gitlab login

mattermost-dl.py

@@ -1,13 +1,19 @@
 import os
+import pickle
 import sqlite3
 from datetime import datetime, date
-from typing import Tuple, Dict, List
+from typing import Tuple, Dict, List, Optional
 import getpass
 
+import requests
+from bs4 import BeautifulSoup
 from mattermostdriver import Driver
 import pathlib
 import json
 
+GITLAB_URL_DEFAULT = "https://mmgit.eltex-co.ru"
+COOKIE_FILE_DEFAULT = "cookies.pkl"
+
 
 def connect(host: str, login_token: str = None, username: str = None, password: str = None) -> Driver:
     d = Driver({
@@ -21,6 +27,117 @@ def connect(host: str, login_token: str = None, username: str = None, password:
     return d
 
 
+def _mattermost_url(host: str) -> str:
+    return f"https://{host}"
+
+
+def _save_cookies(session: requests.Session, filename: str) -> None:
+    try:
+        with open(filename, "wb") as f:
+            pickle.dump(session.cookies, f)
+    except OSError as e:
+        print(f"Warning: failed to save cookies to {filename}: {e}")
+
+
+def _load_cookies(session: requests.Session, filename: str) -> bool:
+    if not os.path.exists(filename):
+        return False
+    try:
+        with open(filename, "rb") as f:
+            session.cookies.update(pickle.load(f))
+        return True
+    except (OSError, pickle.PickleError) as e:
+        print(f"Warning: failed to load cookies from {filename}: {e}")
+    return False
+
+
+def _remove_cookie_file(filename: str) -> None:
+    try:
+        if os.path.exists(filename):
+            os.remove(filename)
+            print(f"Removed invalid cookie file: {filename}")
+    except OSError as e:
+        print(f"Warning: failed to remove cookie file {filename}: {e}")
+
+
+def _get_current_user(session: requests.Session, host: str) -> Optional[dict]:
+    try:
+        response = session.get(f"{_mattermost_url(host)}/api/v4/users/me", timeout=10)
+        return response.json() if response.status_code == 200 else None
+    except requests.exceptions.RequestException:
+        return None
+
+
+def _login_to_mattermost_via_gitlab(session: requests.Session, host: str, gitlab_url: str,
+                                    username: str, password: str) -> bool:
+    login_url = f"{_mattermost_url(host)}/oauth/gitlab/login"
+    headers = {"User-Agent": "Mozilla/5.0"}
+
+    try:
+        response = session.get(login_url, headers=headers, allow_redirects=True, timeout=15)
+        response.raise_for_status()
+
+        if "/users/sign_in" in response.url or "/users/auth/ldapmain" in response.url:
+            soup = BeautifulSoup(response.text, "html.parser")
+            auth_form = soup.find("form", {"id": "new_ldap_user"})
+            if not auth_form:
+                return False
+            auth_token_input = auth_form.find("input", {"name": "authenticity_token"})
+            if not auth_token_input:
+                return False
+
+            auth_token = auth_token_input["value"]
+            payload = {"authenticity_token": auth_token, "username": username, "password": password}
+            callback_url = f"{gitlab_url}/users/auth/ldapmain/callback"
+            response = session.post(
+                callback_url, data=payload, headers={"Referer": response.url},
+                allow_redirects=True, timeout=15
+            )
+            response.raise_for_status()
+
+        if response.status_code == 200 and "/oauth/authorize" in response.url:
+            soup = BeautifulSoup(response.text, "html.parser")
+            redirect_link = soup.find("a", href=True)
+            if not redirect_link:
+                return False
+
+            response = session.get(
+                redirect_link["href"], headers={"Referer": response.url},
+                allow_redirects=True, timeout=15
+            )
+            response.raise_for_status()
+
+    except requests.exceptions.RequestException as e:
+        print(f"GitLab login failed: {e}")
+        return False
+
+    return "MMAUTHTOKEN" in session.cookies
+
+
+def get_token_via_gitlab(host: str, gitlab_url: str, username: str, password: str,
+                         cookie_file: str = COOKIE_FILE_DEFAULT) -> Optional[str]:
+    session = requests.Session()
+
+    if _load_cookies(session, cookie_file):
+        print("Checking saved GitLab/Mattermost session...")
+        if _get_current_user(session, host):
+            token = session.cookies.get("MMAUTHTOKEN")
+            if token:
+                print("Authenticated using saved cookies.")
+                return token
+        print("Saved cookies are invalid or expired.")
+        _remove_cookie_file(cookie_file)
+
+    print("Logging in to Mattermost via GitLab...")
+    if not _login_to_mattermost_via_gitlab(session, host, gitlab_url, username, password):
+        print("Failed to authenticate via GitLab.")
+        return None
+
+    _save_cookies(session, cookie_file)
+    print(f"Session saved to {cookie_file}")
+    return session.cookies.get("MMAUTHTOKEN")
+
+
 def get_users(d: Driver) -> Tuple[Dict[str, str], str]:
     my_user = d.users.get_user("me")
     my_username = my_user["username"]
@@ -216,8 +333,10 @@ def complete_config(config: dict, config_filename: str = "config.json") -> dict:
         print(f"Using login mode '{config['login_mode']}' from config")
     else:
         login_mode = ""
-        while login_mode not in ["password", "token"]:
-            login_mode = input("Please input login_mode 'password' or 'token' (=Gitlab Oauth): ")
+        while login_mode not in ["password", "token", "gitlab"]:
+            login_mode = input(
+                "Please input login_mode 'password', 'token', or 'gitlab' (GitLab OAuth): "
+            )
         config["login_mode"] = login_mode
         config_changed = True
 
@@ -230,20 +349,66 @@ def complete_config(config: dict, config_filename: str = "config.json") -> dict:
             config_changed = True
 
         password = getpass.getpass("Enter your password (hidden): ")
+    elif config["login_mode"] == "gitlab":
+        if config.get("gitlab_url", False):
+            print(f"Using GitLab URL '{config['gitlab_url']}' from config")
+        else:
+            config["gitlab_url"] = input(
+                f"Please input GitLab URL [{GITLAB_URL_DEFAULT}]: "
+            ) or GITLAB_URL_DEFAULT
+            config_changed = True
+
+        if config.get("username", False):
+            print(f"Using username '{config['username']}' from config")
+        else:
+            config["username"] = input("Please input your GitLab username: ")
+            config_changed = True
+
+        password = getpass.getpass("Enter your GitLab password (hidden): ")
+        token = get_token_via_gitlab(
+            config["host"], config["gitlab_url"], config["username"], password,
+            config.get("cookie_file", COOKIE_FILE_DEFAULT)
+        )
+        if not token:
+            raise SystemExit("GitLab authentication failed.")
+        config["token"] = token
+        config_changed = True
     else:
         if config.get("token", False):
             print(f"Using token '{config['token']}' from config")
         else:
-            print("Are you logged-in into Mattermost using the Firefox Browser? "
-                  "If so, token may be automatically extracted")
+            print("Token not found in config. Choose how to obtain MMAUTHTOKEN:")
+            print("  g - login via GitLab (username/password)")
+            print("  f - extract from Firefox cookies")
+            print("  m - enter token manually")
             dec = ""
-            while not (dec == "y" or dec == "n"):
-                dec = input("Try to find token automatically? y/n: ")
+            while dec not in ["g", "f", "m"]:
+                dec = input("Choice [g/f/m]: ")
 
             token = None
-            if dec == "y":
+            if dec == "g":
+                if config.get("gitlab_url", False):
+                    print(f"Using GitLab URL '{config['gitlab_url']}' from config")
+                else:
+                    config["gitlab_url"] = input(
+                        f"Please input GitLab URL [{GITLAB_URL_DEFAULT}]: "
+                    ) or GITLAB_URL_DEFAULT
+                    config_changed = True
+                if config.get("username", False):
+                    print(f"Using username '{config['username']}' from config")
+                else:
+                    config["username"] = input("Please input your GitLab username: ")
+                    config_changed = True
+                password = getpass.getpass("Enter your GitLab password (hidden): ")
+                token = get_token_via_gitlab(
+                    config["host"], config["gitlab_url"], config["username"], password,
+                    config.get("cookie_file", COOKIE_FILE_DEFAULT)
+                )
+                if not token:
+                    raise SystemExit("GitLab authentication failed.")
+            elif dec == "f":
                 token = find_mmauthtoken_firefox(config["host"])
-            elif not token:
+            if not token:
                 token = input("Please input your login token (MMAUTHTOKEN): ")
             config["token"] = token
             config_changed = True