Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
auth_strategy.py307 linesDownload Raw Back to paramiko
1"""2Modern, adaptable authentication machinery.3 4Replaces certain parts of `.SSHClient`. For a concrete implementation, see the5``OpenSSHAuthStrategy`` class in `Fabric <https://fabfile.org>`_.6"""7 8from collections import namedtuple9 10from .agent import AgentKey11from .util import get_logger12from .ssh_exception import AuthenticationException13 14 15class AuthSource:16    """17    Some SSH authentication source, such as a password, private key, or agent.18 19    See subclasses in this module for concrete implementations.20 21    All implementations must accept at least a ``username`` (``str``) kwarg.22    """23 24    def __init__(self, username):25        self.username = username26 27    def _repr(self, **kwargs):28        # TODO: are there any good libs for this? maybe some helper from29        # structlog?30        pairs = [f"{k}={v!r}" for k, v in kwargs.items()]31        joined = ", ".join(pairs)32        return f"{self.__class__.__name__}({joined})"33 34    def __repr__(self):35        return self._repr()36 37    def authenticate(self, transport):38        """39        Perform authentication.40        """41        raise NotImplementedError42 43 44class NoneAuth(AuthSource):45    """46    Auth type "none", ie https://www.rfc-editor.org/rfc/rfc4252#section-5.2 .47    """48 49    def authenticate(self, transport):50        return transport.auth_none(self.username)51 52 53class Password(AuthSource):54    """55    Password authentication.56 57    :param callable password_getter:58        A lazy callable that should return a `str` password value at59        authentication time, such as a `functools.partial` wrapping60        `getpass.getpass`, an API call to a secrets store, or similar.61 62        If you already know the password at instantiation time, you should63        simply use something like ``lambda: "my literal"`` (for a literal, but64        also, shame on you!) or ``lambda: variable_name`` (for something stored65        in a variable).66    """67 68    def __init__(self, username, password_getter):69        super().__init__(username=username)70        self.password_getter = password_getter71 72    def __repr__(self):73        # Password auth is marginally more 'username-caring' than pkeys, so may74        # as well log that info here.75        return super()._repr(user=self.username)76 77    def authenticate(self, transport):78        # Lazily get the password, in case it's prompting a user79        # TODO: be nice to log source _of_ the password?80        password = self.password_getter()81        return transport.auth_password(self.username, password)82 83 84# TODO 4.0: twiddle this, or PKey, or both, so they're more obviously distinct.85# TODO 4.0: the obvious is to make this more wordy (PrivateKeyAuth), the86# minimalist approach might be to rename PKey to just Key (esp given all the87# subclasses are WhateverKey and not WhateverPKey)88class PrivateKey(AuthSource):89    """90    Essentially a mixin for private keys.91 92    Knows how to auth, but leaves key material discovery/loading/decryption to93    subclasses.94 95    Subclasses **must** ensure that they've set ``self.pkey`` to a decrypted96    `.PKey` instance before calling ``super().authenticate``; typically97    either in their ``__init__``, or in an overridden ``authenticate`` prior to98    its `super` call.99    """100 101    def authenticate(self, transport):102        return transport.auth_publickey(self.username, self.pkey)103 104 105class InMemoryPrivateKey(PrivateKey):106    """107    An in-memory, decrypted `.PKey` object.108    """109 110    def __init__(self, username, pkey):111        super().__init__(username=username)112        # No decryption (presumably) necessary!113        self.pkey = pkey114 115    def __repr__(self):116        # NOTE: most of interesting repr-bits for private keys is in PKey.117        # TODO: tacking on agent-ness like this is a bit awkward, but, eh?118        rep = super()._repr(pkey=self.pkey)119        if isinstance(self.pkey, AgentKey):120            rep += " [agent]"121        return rep122 123 124class OnDiskPrivateKey(PrivateKey):125    """126    Some on-disk private key that needs opening and possibly decrypting.127 128    :param str source:129        String tracking where this key's path was specified; should be one of130        ``"ssh-config"``, ``"python-config"``, or ``"implicit-home"``.131    :param Path path:132        The filesystem path this key was loaded from.133    :param PKey pkey:134        The `PKey` object this auth source uses/represents.135    """136 137    def __init__(self, username, source, path, pkey):138        super().__init__(username=username)139        self.source = source140        allowed = ("ssh-config", "python-config", "implicit-home")141        if source not in allowed:142            raise ValueError(f"source argument must be one of: {allowed!r}")143        self.path = path144        # Superclass wants .pkey, other two are mostly for display/debugging.145        self.pkey = pkey146 147    def __repr__(self):148        return self._repr(149            key=self.pkey, source=self.source, path=str(self.path)150        )151 152 153# TODO re sources: is there anything in an OpenSSH config file that doesn't fit154# into what Paramiko already had kwargs for?155 156 157SourceResult = namedtuple("SourceResult", ["source", "result"])158 159# TODO: tempting to make this an OrderedDict, except the keys essentially want160# to be rich objects (AuthSources) which do not make for useful user indexing?161# TODO: members being vanilla tuples is pretty old-school/expedient; they162# "really" want to be something that's type friendlier (unless the tuple's 2nd163# member being a Union of two types is "fine"?), which I assume means yet more164# classes, eg an abstract SourceResult with concrete AuthSuccess and165# AuthFailure children?166# TODO: arguably we want __init__ typechecking of the members (or to leverage167# mypy by classifying this literally as list-of-AuthSource?)168class AuthResult(list):169    """170    Represents a partial or complete SSH authentication attempt.171 172    This class conceptually extends `AuthStrategy` by pairing the former's173    authentication **sources** with the **results** of trying to authenticate174    with them.175 176    `AuthResult` is a (subclass of) `list` of `namedtuple`, which are of the177    form ``namedtuple('SourceResult', 'source', 'result')`` (where the178    ``source`` member is an `AuthSource` and the ``result`` member is either a179    return value from the relevant `.Transport` method, or an exception180    object).181 182    .. note::183        Transport auth method results are always themselves a ``list`` of "next184        allowable authentication methods".185 186        In the simple case of "you just authenticated successfully", it's an187        empty list; if your auth was rejected but you're allowed to try again,188        it will be a list of string method names like ``pubkey`` or189        ``password``.190 191        The ``__str__`` of this class represents the empty-list scenario as the192        word ``success``, which should make reading the result of an193        authentication session more obvious to humans.194 195    Instances also have a `strategy` attribute referencing the `AuthStrategy`196    which was attempted.197    """198 199    def __init__(self, strategy, *args, **kwargs):200        self.strategy = strategy201        super().__init__(*args, **kwargs)202 203    def __str__(self):204        # NOTE: meaningfully distinct from __repr__, which still wants to use205        # superclass' implementation.206        # TODO: go hog wild, use rich.Table? how is that on degraded term's?207        # TODO: test this lol208        return "\n".join(209            f"{x.source} -> {x.result or 'success'}" for x in self210        )211 212 213# TODO 4.0: descend from SSHException or even just Exception214class AuthFailure(AuthenticationException):215    """216    Basic exception wrapping an `AuthResult` indicating overall auth failure.217 218    Note that `AuthFailure` descends from `AuthenticationException` but is219    generally "higher level"; the latter is now only raised by individual220    `AuthSource` attempts and should typically only be seen by users when221    encapsulated in this class. It subclasses `AuthenticationException`222    primarily for backwards compatibility reasons.223    """224 225    def __init__(self, result):226        self.result = result227 228    def __str__(self):229        return "\n" + str(self.result)230 231 232class AuthStrategy:233    """234    This class represents one or more attempts to auth with an SSH server.235 236    By default, subclasses must at least accept an ``ssh_config``237    (`.SSHConfig`) keyword argument, but may opt to accept more as needed for238    their particular strategy.239    """240 241    def __init__(242        self,243        ssh_config,244    ):245        self.ssh_config = ssh_config246        self.log = get_logger(__name__)247 248    def get_sources(self):249        """250        Generator yielding `AuthSource` instances, in the order to try.251 252        This is the primary override point for subclasses: you figure out what253        sources you need, and ``yield`` them.254 255        Subclasses _of_ subclasses may find themselves wanting to do things256        like filtering or discarding around a call to `super`.257        """258        raise NotImplementedError259 260    def authenticate(self, transport):261        """262        Handles attempting `AuthSource` instances yielded from `get_sources`.263 264        You *normally* won't need to override this, but it's an option for265        advanced users.266        """267        succeeded = False268        overall_result = AuthResult(strategy=self)269        # TODO: arguably we could fit in a "send none auth, record allowed auth270        # types sent back" thing here as OpenSSH-client does, but that likely271        # wants to live in fabric.OpenSSHAuthStrategy as not all target servers272        # will implement it!273        # TODO: needs better "server told us too many attempts" checking!274        for source in self.get_sources():275            self.log.debug(f"Trying {source}")276            try:  # NOTE: this really wants to _only_ wrap the authenticate()!277                result = source.authenticate(transport)278                succeeded = True279            # TODO: 'except PartialAuthentication' is needed for 2FA and280            # similar, as per old SSHClient.connect - it is the only way281            # AuthHandler supplies access to the 'name-list' field from282            # MSG_USERAUTH_FAILURE, at present.283            except Exception as e:284                result = e285                # TODO: look at what this could possibly raise, we don't really286                # want Exception here, right? just SSHException subclasses? or287                # do we truly want to capture anything at all with assumption288                # it's easy enough for users to look afterwards?289                # NOTE: showing type, not message, for tersity & also most of290                # the time it's basically just "Authentication failed."291                source_class = e.__class__.__name__292                self.log.info(293                    f"Authentication via {source} failed with {source_class}"294                )295            overall_result.append(SourceResult(source, result))296            if succeeded:297                break298        # Gotta die here if nothing worked, otherwise Transport's main loop299        # just kinda hangs out until something times out!300        if not succeeded:301            raise AuthFailure(result=overall_result)302        # Success: give back what was done, in case they care.303        return overall_result304 305    # TODO: is there anything OpenSSH client does which _can't_ cleanly map to306    # iterating a generator?307