# Copyright (C) 2016-2024 RhodeCode GmbH
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License, version 3
# (only), as published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# This program is dual-licensed. If you wish to learn more about the
# RhodeCode Enterprise Edition, including its added features, Support services,
# and proprietary license terms, please see https://rhodecode.com/licenses/

import os
import re
import logging
import datetime
from sqlalchemy import Table

from rhodecode.lib.api_utils import call_service_api
from rhodecode.lib.utils2 import AttributeDict

from .hg import MercurialServer
from .git import GitServer
from .svn import SubversionServer

log = logging.getLogger(__name__)


class SshWrapper(object):
    hg_cmd_pat = re.compile(r"^hg\s+\-R\s+(\S+)\s+serve\s+\-\-stdio$")
    git_cmd_pat = re.compile(r"^git-(receive-pack|upload-pack)\s\'[/]?(\S+?)(|\.git)\'$")
    svn_cmd_pat = re.compile(r"^svnserve -t")

    def __init__(self, command, connection_info, mode, user, user_id, key_id: int, shell, ini_path: str, settings, env):
        self.command = command
        self.connection_info = connection_info
        self.mode = mode
        self.username = user
        self.user_id = user_id
        self.key_id = key_id
        self.shell = shell
        self.ini_path = ini_path
        self.env = env
        self.settings = settings
        self.server_impl = None

    def update_key_access_time(self, key_id):
        from rhodecode.model.meta import raw_query_executor, Base

        table = Table("user_ssh_keys", Base.metadata, autoload=False)
        atime = datetime.datetime.utcnow()
        stmt = (
            table.update().where(table.c.ssh_key_id == key_id).values(accessed_on=atime)
            # no MySQL Support for .returning :((
            # .returning(table.c.accessed_on, table.c.ssh_key_fingerprint)
        )

        res_count = None
        with raw_query_executor() as session:
            result = session.execute(stmt)
            if result.rowcount:
                res_count = result.rowcount

        if res_count:
            log.debug("Update key id:`%s` access time", key_id)

    def get_user(self, user_id):
        user = AttributeDict()
        # lazy load db imports
        from rhodecode.model.db import User

        dbuser = User.get(user_id)
        if not dbuser:
            return None
        user.user_id = dbuser.user_id
        user.username = dbuser.username
        user.auth_user = dbuser.AuthUser()
        return user

    def get_connection_info(self):
        """
        connection_info

        Identifies the client and server ends of the connection.
        The variable contains four space-separated values: client IP address,
        client port number, server IP address, and server port number.
        """
        conn = dict(
            client_ip=None,
            client_port=None,
            server_ip=None,
            server_port=None,
        )

        info = self.connection_info.split(" ")
        if len(info) == 4:
            conn["client_ip"] = info[0]
            conn["client_port"] = info[1]
            conn["server_ip"] = info[2]
            conn["server_port"] = info[3]

        return conn

    def maybe_translate_repo_uid(self, repo_name):
        _org_name = repo_name
        if _org_name.startswith("_"):
            # remove format of _ID/subrepo
            _org_name = _org_name.split("/", 1)[0]

        if repo_name.startswith("_"):
            from rhodecode.model.repo import RepoModel

            org_repo_name = repo_name
            log.debug("translating UID repo %s", org_repo_name)
            by_id_match = RepoModel().get_repo_by_id(repo_name)
            if by_id_match:
                repo_name = by_id_match.repo_name
                log.debug("translation of UID repo %s got `%s`", org_repo_name, repo_name)

        return repo_name, _org_name

    def get_repo_details(self, mode):
        vcs_type = mode if mode in ["svn", "hg", "git"] else None
        repo_name = None

        hg_match = self.hg_cmd_pat.match(self.command)
        if hg_match is not None:
            vcs_type = "hg"
            repo_id = hg_match.group(1).strip("/")
            repo_name, org_name = self.maybe_translate_repo_uid(repo_id)
            return vcs_type, repo_name, mode

        git_match = self.git_cmd_pat.match(self.command)
        if git_match is not None:
            mode = git_match.group(1)
            vcs_type = "git"
            repo_id = git_match.group(2).strip("/")
            repo_name, org_name = self.maybe_translate_repo_uid(repo_id)
            return vcs_type, repo_name, mode

        svn_match = self.svn_cmd_pat.match(self.command)
        if svn_match is not None:
            vcs_type = "svn"
            # Repo name should be extracted from the input stream, we're unable to
            # extract it at this point in execution
            return vcs_type, repo_name, mode

        return vcs_type, repo_name, mode

    def serve(self, vcs, repo, mode, user, permissions, branch_permissions):
        # TODO: remove this once we have .ini defined access path...
        from rhodecode.model.scm import ScmModel

        store = ScmModel().repos_path

        check_branch_perms = False
        detect_force_push = False

        if branch_permissions:
            check_branch_perms = True
            detect_force_push = True

        log.debug(
            "VCS detected:`%s` mode: `%s` repo_name: %s, branch_permission_checks:%s",
            vcs,
            mode,
            repo,
            check_branch_perms,
        )

        # detect if we have to check branch permissions
        extras = {
            "detect_force_push": detect_force_push,
            "check_branch_perms": check_branch_perms,
            "config": self.ini_path,
        }

        if vcs == "hg":
            server = MercurialServer(
                store=store,
                ini_path=self.ini_path,
                repo_name=repo,
                user=user,
                user_permissions=permissions,
                settings=self.settings,
                env=self.env,
            )
            self.server_impl = server
            return server.run(tunnel_extras=extras)

        elif vcs == "git":
            server = GitServer(
                store=store,
                ini_path=self.ini_path,
                repo_name=repo,
                repo_mode=mode,
                user=user,
                user_permissions=permissions,
                settings=self.settings,
                env=self.env,
            )
            self.server_impl = server
            return server.run(tunnel_extras=extras)

        elif vcs == "svn":
            server = SubversionServer(
                store=store,
                ini_path=self.ini_path,
                repo_name=None,
                user=user,
                user_permissions=permissions,
                settings=self.settings,
                env=self.env,
            )
            self.server_impl = server
            return server.run(tunnel_extras=extras)

        else:
            raise Exception(f"Unrecognised VCS: {vcs}")

    def wrap(self):
        mode = self.mode
        username = self.username
        user_id = self.user_id
        key_id = self.key_id
        shell = self.shell

        scm_detected, scm_repo, scm_mode = self.get_repo_details(mode)

        log.debug(
            ""
            "mode:`%s` "
            "user_name:`%s`"
            "user_id:`%s`"
            "shell: `%s` "
            'ssh_command:`"%s"` '
            "scm_mode:`%s` "
            "scm_type:`%s` "
            "scm_repo_name:`%s`",
            mode,
            username,
            user_id,
            shell,
            self.command,
            scm_mode,
            scm_detected,
            scm_repo,
        )

        log.debug("SSH Connection info %s", self.get_connection_info())

        # update last access time for this key
        if key_id:
            self.update_key_access_time(key_id)

        if shell and self.command is None:
            log.info("Dropping to shell, no command given and shell is allowed")
            os.execl("/bin/bash", "-l")
            exit_code = 1

        elif scm_detected:
            user = self.get_user(user_id)
            if not user:
                log.warning("User with id %s not found", user_id)
                exit_code = -1
                return exit_code

            auth_user = user.auth_user
            permissions = auth_user.permissions["repositories"]
            repo_branch_permissions = auth_user.get_branch_permissions(scm_repo)
            try:
                exit_code, is_updated = self.serve(
                    scm_detected, scm_repo, scm_mode, user, permissions, repo_branch_permissions
                )
            except Exception:
                log.exception("Error occurred during execution of SshWrapper")
                exit_code = -1

        elif self.command is None and shell is False:
            log.error("No Command given.")
            exit_code = -1

        else:
            log.error('Unhandled Command: "%s" Aborting.', self.command)
            exit_code = -1

        return exit_code


class SshWrapperStandalone(SshWrapper):
    """
    New version of SshWrapper designed to be depended only on service API
    """

    repos_path = None

    @staticmethod
    def parse_user_related_data(user_data):
        user = AttributeDict()
        user.user_id = user_data["user_id"]
        user.username = user_data["username"]
        user.repo_permissions = user_data["repo_permissions"]
        user.branch_permissions = user_data["branch_permissions"]
        return user

    def wrap(self):
        mode = self.mode
        username = self.username
        user_id = self.user_id
        shell = self.shell

        scm_detected, scm_repo, scm_mode = self.get_repo_details(mode)

        log.debug(
            'Mode: `%s` User: `name:%s : id:%s` Shell: `%s` SSH Command: `"%s"` '
            "SCM_DETECTED: `%s` SCM Mode: `%s` SCM Repo: `%s`",
            mode,
            username,
            user_id,
            shell,
            self.command,
            scm_detected,
            scm_mode,
            scm_repo,
        )

        log.debug("SSH Connection info %s", self.get_connection_info())

        if shell and self.command is None:
            log.info("Dropping to shell, no command given and shell is allowed")
            os.execl("/bin/bash", "-l")
            exit_code = 1

        elif scm_detected:
            data = call_service_api(
                self.settings,
                {
                    "method": "service_get_data_for_ssh_wrapper",
                    "args": {"user_id": user_id, "repo_name": scm_repo, "key_id": self.key_id},
                },
            )
            if not data:
                log.error("Error occurred during execution of SshWrapper service API call")

            user = self.parse_user_related_data(data)
            if not user:
                log.warning("User with id %s not found", user_id)
                exit_code = -1
                return exit_code
            self.repos_path = data["repos_path"]
            permissions = user.repo_permissions
            repo_branch_permissions = user.branch_permissions
            try:
                exit_code, is_updated = self.serve(
                    scm_detected, scm_repo, scm_mode, user, permissions, repo_branch_permissions
                )
            except Exception:
                log.exception("Error occurred during execution of SshWrapper")
                exit_code = -1

        elif self.command is None and shell is False:
            log.error("No Command given.")
            exit_code = -1

        else:
            log.error('Unhandled Command: "%s" Aborting.', self.command)
            exit_code = -1

        return exit_code

    def maybe_translate_repo_uid(self, repo_name):
        _org_name = repo_name
        if _org_name.startswith("_"):
            _org_name = _org_name.split("/", 1)[0]

        if repo_name.startswith("_"):
            org_repo_name = repo_name
            log.debug("translating UID repo %s", org_repo_name)
            by_id_match = call_service_api(
                self.settings, {"method": "service_get_repo_name_by_id", "args": {"repo_id": repo_name}}
            )
            if by_id_match:
                repo_name = by_id_match["repo_name"]
                log.debug("translation of UID repo %s got `%s`", org_repo_name, repo_name)

        return repo_name, _org_name

    def serve(self, vcs, repo, mode, user, permissions, branch_permissions):
        store = self.repos_path

        check_branch_perms = False
        detect_force_push = False

        if branch_permissions:
            check_branch_perms = True
            detect_force_push = True

        log.debug(
            "VCS detected:`%s` mode: `%s` repo_name: %s, branch_permission_checks:%s",
            vcs,
            mode,
            repo,
            check_branch_perms,
        )

        # detect if we have to check branch permissions
        extras = {
            "detect_force_push": detect_force_push,
            "check_branch_perms": check_branch_perms,
            "config": self.ini_path,
        }

        match vcs:
            case "hg":
                server = MercurialServer(
                    store=store,
                    ini_path=self.ini_path,
                    repo_name=repo,
                    user=user,
                    user_permissions=permissions,
                    settings=self.settings,
                    env=self.env,
                )
            case "git":
                server = GitServer(
                    store=store,
                    ini_path=self.ini_path,
                    repo_name=repo,
                    repo_mode=mode,
                    user=user,
                    user_permissions=permissions,
                    settings=self.settings,
                    env=self.env,
                )
            case "svn":
                server = SubversionServer(
                    store=store,
                    ini_path=self.ini_path,
                    repo_name=None,
                    user=user,
                    user_permissions=permissions,
                    settings=self.settings,
                    env=self.env,
                )
            case _:
                raise Exception(f"Unrecognised VCS: {vcs}")
        self.server_impl = server
        return server.run(tunnel_extras=extras)
