From c8fd9822eae845a913b4da39607693f780e70379 Mon Sep 17 00:00:00 2001
From: Dinis <dlage@users.noreply.github.com>
Date: Mon, 1 Apr 2019 00:36:39 +0100
Subject: [PATCH] Update tmux.py

---
 scripts/tmux.py | 70 +++++++++++++++++++++++++++++++++----------------
 1 file changed, 48 insertions(+), 22 deletions(-)

diff --git a/scripts/tmux.py b/scripts/tmux.py
index 3cc70b3..6e7fea5 100755
--- a/scripts/tmux.py
+++ b/scripts/tmux.py
@@ -1,28 +1,54 @@
 #!/usr/bin/env python
 import os
 import subprocess
+import sys
 
-print("TMUX: %s" % os.getenv("TMUX", ""))
-print("SSH_TTY: %s" % os.getenv("SSH_TTY", ""))
-print("SSH_AUTH_SOCK: %s" % os.getenv("SSH_AUTH_SOCK", ""))
-print("HOME: %s" % os.getenv("HOME", ""))
+session_name  = 'sshwrap'
+if len(sys.argv) > 1:
+    session_name = sys.argv[1]
 
-if os.getenv("TMUX") is None:
-    if os.getenv("SSH_TTY") is not None:
-        if os.getenv("SSH_AUTH_SOCK") is not None:
-            sock_file = os.path.join(os.getenv("HOME"), ".wrap_auth_sock")
-            try:
-                #always try to remove
-                os.remove(sock_file)
-            except OSError:
-                pass
-            os.symlink(os.getenv("SSH_AUTH_SOCK"), sock_file)
-            os.environ['SSH_AUTH_SOCK'] = sock_file
+ssh_auth_sock = '-'.join(['ssh_auth_sock', os.getenv('HOSTNAME'), session_name])
 
-    try:
-        subprocess.check_call(["tmux", "attach-session", "-t", "sshwrap"])
-    except subprocess.CalledProcessError:
-        print("lets create session")
-        os.environ['STY'] = "tmux-sshwrap"
-        os.execlpe("tmux", "tmux", "new-session", "-s", "sshwrap",
-                   os.environ)
+def is_valid(key):
+    if not key.startswith('PYENV'):
+        return True
+
+    if key.endswith('_ROOT'):
+        return True
+
+    if key.endswith('_SHELL'):
+        return True
+
+    return False
+
+
+def main():
+    if os.getenv("TMUX") is None:
+        if os.getenv("SSH_TTY") is not None:
+            if os.getenv("SSH_AUTH_SOCK") is not None:
+                sock_file = os.path.join(os.getenv("HOME"), ".ssh", ssh_auth_sock)
+                try:
+                    #always try to remove
+                    os.remove(sock_file)
+                except OSError:
+                    pass
+                os.symlink(os.getenv("SSH_AUTH_SOCK"), sock_file)
+                os.environ['SSH_AUTH_SOCK'] = sock_file
+
+        try:
+            subprocess.check_call(["tmux", "attach-session", "-t", session_name])
+        except subprocess.CalledProcessError:
+            os.environ['STY'] = "tmux-sshwrap"
+            env = {k: os.environ[k] for k in os.environ.keys() if is_valid(k)}
+            os.execlpe("tmux", "tmux", "new-session", "-s", session_name, env)
+    else:
+        print('TMUX var is set...')
+
+
+if __name__ == '__main__':
+    print("TMUX: %s" % os.getenv("TMUX", ""))
+    print("SSH_TTY: %s" % os.getenv("SSH_TTY", ""))
+    print("SSH_AUTH_SOCK: %s" % os.getenv("SSH_AUTH_SOCK", ""))
+    print("HOME: %s" % os.getenv("HOME", ""))
+
+    main()