diff --git a/mattermost-dl.py b/mattermost-dl.py index ac03bb7..69a167a 100644 --- a/mattermost-dl.py +++ b/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