<commit_before>'''
This file as originally part of the "sshuserclient" pypi package. The GitHub
source has now vanished (https://github.com/tobald/sshuserclient).
'''

import os

from paramiko import (
    AutoAddPolicy,
    ProxyCommand,
    SSHClient as ParamikoClient,
)

from .config import SSHConfig


class SSHClient(ParamikoClient):
    '''
    An SSHClient which honors ssh_config and supports proxyjumping
    original idea at http://bitprophet.org/blog/2012/11/05/gateway-solutions/
    '''

    def connect(self, hostname, **kwargs):
        self.hostname, self.config = self.parse_config(hostname)
        self.config.update(kwargs)
        super(SSHClient, self).connect(hostname, **self.config)

    def gateway(self, target, target_port):
        transport = self.get_transport()
        return transport.open_channel(
            'direct-tcpip',
            (target, target_port),
            (self.hostname, self.config['port']))

    def parse_config(self, hostname):
        cfg = {'port': 22}

        user_config_file = os.path.expanduser('~/.ssh/config')
        if os.path.exists(user_config_file):
            with open(user_config_file) as f:
                ssh_config = SSHConfig()
                ssh_config.parse(f)
                host_config = ssh_config.lookup(hostname)
                if 'hostname' in host_config:
                    hostname = host_config['hostname']
                if 'user' in host_config:
                    cfg['username'] = host_config['user']
                if 'identityfile' in host_config:
                    cfg['key_filename'] = host_config['identityfile']
                if 'port' in host_config:
                    cfg['port'] = int(host_config['port'])
                if 'proxycommand' in host_config:
                    cfg['sock'] = ProxyCommand(host_config['proxycommand'])
                elif 'proxyjump' in host_config:
                    hops = host_config['proxyjump'].split(',')
                    sock = None
                    for i, hop in enumerate(hops):
                        hop_hostname, hop_config = self.derive_shorthand(hop)
                        c = SSHClient()
                        c.set_missing_host_key_policy(AutoAddPolicy())
                        c.connect(hop_hostname, sock=sock, **hop_config)
                        if i == len(hops) - 1:
                            target = hostname
                            target_config = {'port': cfg['port']}
                        else:
                            target, target_config = self.derive_shorthand(
                                hops[i + 1])
                        sock = c.gateway(target, target_config['port'])
                    cfg['sock'] = sock
        return hostname, cfg

    def derive_shorthand(self, host_string):
        config = {}
        user_hostport = host_string.rsplit('@', 1)
        hostport = user_hostport.pop()
        user = user_hostport[0] if user_hostport and user_hostport[0] else None
        if user:
            config['username'] = user

        # IPv6: can't reliably tell where addr ends and port begins, so don't
        # try (and don't bother adding special syntax either, user should avoid<commit_msg>Only parse user SSH config once.<commit_after>'''
This file as originally part of the "sshuserclient" pypi package. The GitHub
source has now vanished (https://github.com/tobald/sshuserclient).
'''

import os

from paramiko import (
    AutoAddPolicy,
    ProxyCommand,
    SSHClient as ParamikoClient,
)

from pyinfra.api.util import memoize

from .config import SSHConfig


@memoize
def get_ssh_config():
    user_config_file = os.path.expanduser('~/.ssh/config')
    if os.path.exists(user_config_file):
        with open(user_config_file) as f:
            ssh_config = SSHConfig()
            ssh_config.parse(f)
            return ssh_config


class SSHClient(ParamikoClient):
    '''
    An SSHClient which honors ssh_config and supports proxyjumping
    original idea at http://bitprophet.org/blog/2012/11/05/gateway-solutions/
    '''

    def connect(self, hostname, **kwargs):
        self.hostname, self.config = self.parse_config(hostname)
        self.config.update(kwargs)
        super(SSHClient, self).connect(hostname, **self.config)

    def gateway(self, target, target_port):
        transport = self.get_transport()
        return transport.open_channel(
            'direct-tcpip',
            (target, target_port),
            (self.hostname, self.config['port']))

    def parse_config(self, hostname):
        cfg = {'port': 22}

        ssh_config = get_ssh_config()
        host_config = ssh_config.lookup(hostname)

        if 'hostname' in host_config:
            hostname = host_config['hostname']
        if 'user' in host_config:
            cfg['username'] = host_config['user']
        if 'identityfile' in host_config:
            cfg['key_filename'] = host_config['identityfile']
        if 'port' in host_config:
            cfg['port'] = int(host_config['port'])
        if 'proxycommand' in host_config:
            cfg['sock'] = ProxyCommand(host_config['proxycommand'])
        elif 'proxyjump' in host_config:
            hops = host_config['proxyjump'].split(',')
            sock = None
            for i, hop in enumerate(hops):
                hop_hostname, hop_config = self.derive_shorthand(hop)
                c = SSHClient()
                c.set_missing_host_key_policy(AutoAddPolicy())
                c.connect(hop_hostname, sock=sock, **hop_config)
                if i == len(hops) - 1:
                    target = hostname
                    target_config = {'port': cfg['port']}
                else:
                    target, target_config = self.derive_shorthand(
                        hops[i + 1])
                sock = c.gateway(target, target_config['port'])
            cfg['sock'] = sock

        return hostname, cfg

    def derive_shorthand(self, host_string):
        config = {}
        user_hostport = host_string.rsplit('@', 1)
        hostport = user_hostport.pop()
        user = user_hostport[0] if user_hostport and user_hostport[0] else None
        if user:
            config['username'] = user

        # IPv6: can't reliably tell where addr ends and port begins, so don't
        # try (and don't bother adding special syntax either, user should avoid
{'fraction_non_alphanumeric': 0.10314465408805032, 'fraction_numerical': 0.005031446540880503, 'mean_word_length': 2.0377268385864373, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5146258', 'n_tokens_mistral': 1856, 'n_tokens_neox': 1739, 'n_words': 431}