select-session.py 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209
  1. #! /usr/bin/python3
  2. #
  3. # select-session.py
  4. # Copyright (C) 2010 Canonical Ltd.
  5. # Copyright (C) 2012-2014 Dustin Kirkland <kirkland@byobu.org>
  6. #
  7. # Authors: Dustin Kirkland <kirkland@byobu.org>
  8. # Ryan C. Thompson <rct@thompsonclan.org>
  9. #
  10. # This program is free software: you can redistribute it and/or modify
  11. # it under the terms of the GNU General Public License as published by
  12. # the Free Software Foundation, version 3 of the License.
  13. #
  14. # This program is distributed in the hope that it will be useful,
  15. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  16. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  17. # GNU General Public License for more details.
  18. #
  19. # You should have received a copy of the GNU General Public License
  20. # along with this program. If not, see <http://www.gnu.org/licenses/>.
  21. import os
  22. import re
  23. import sys
  24. import subprocess
  25. try:
  26. # For Python3, try and import input from builtins
  27. from builtins import input
  28. except Exception:
  29. # But fall back to using the default input
  30. True
  31. PKG = "byobu"
  32. SHELL = os.getenv("SHELL", "/bin/bash")
  33. HOME = os.getenv("HOME")
  34. BYOBU_CONFIG_DIR = os.getenv("BYOBU_CONFIG_DIR", HOME + "/.byobu")
  35. BYOBU_BACKEND = os.getenv("BYOBU_BACKEND", "tmux")
  36. choice = -1
  37. sessions = []
  38. text = []
  39. reuse_sessions = os.path.exists("%s/.reuse-session" % (BYOBU_CONFIG_DIR))
  40. BYOBU_UPDATE_ENVVARS = ["DISPLAY", "DBUS_SESSION_BUS_ADDRESS", "SESSION_MANAGER", "GPG_AGENT_INFO", "XDG_SESSION_COOKIE", "XDG_SESSION_PATH", "GNOME_KEYRING_CONTROL", "GNOME_KEYRING_PID", "GPG_AGENT_INFO", "SSH_ASKPASS", "SSH_AUTH_SOCK", "SSH_AGENT_PID", "WINDOWID", "UPSTART_JOB", "UPSTART_EVENTS", "UPSTART_SESSION", "UPSTART_INSTANCE"]
  41. def get_sessions():
  42. sessions = []
  43. i = 0
  44. output = False
  45. if BYOBU_BACKEND == "screen":
  46. try:
  47. output = subprocess.Popen(["screen", "-ls"], stdout=subprocess.PIPE).communicate()[0]
  48. except subprocess.CalledProcessError as cpe:
  49. # screen -ls seems to always return 1
  50. if cpe.returncode != 1:
  51. raise
  52. else:
  53. output = cpe.output
  54. if sys.stdout.encoding is None:
  55. output = output.decode("UTF-8")
  56. else:
  57. output = output.decode(sys.stdout.encoding)
  58. if output:
  59. for s in output.splitlines():
  60. s = re.sub(r'\s+', ' ', s)
  61. # Ignore hidden sessions (named sessions that start with a "." or a "_")
  62. if s and s != " " and (s.find(" ") == 0 and len(s) > 1 and s.count("..") == 0 and s.count("._") == 0):
  63. text.append("screen: %s" % s.strip())
  64. items = s.split(" ")
  65. sessions.append("screen____%s" % items[1])
  66. i += 1
  67. if BYOBU_BACKEND == "tmux":
  68. output = subprocess.Popen(["tmux", "list-sessions"], stdout=subprocess.PIPE).communicate()[0]
  69. if sys.stdout.encoding is None:
  70. output = output.decode("UTF-8")
  71. else:
  72. output = output.decode(sys.stdout.encoding)
  73. if output:
  74. for s in output.splitlines():
  75. # Ignore hidden sessions (named sessions that start with a "_")
  76. if s and not s.startswith("_") and s.find("-") == -1:
  77. text.append("tmux: %s" % s.strip())
  78. sessions.append("tmux____%s" % s.split(":")[0])
  79. i += 1
  80. return sessions
  81. def cull_zombies(session_name):
  82. # When using tmux session groups, closing a client will leave
  83. # unattached "zombie" sessions that will never be reattached.
  84. # Search for and kill any unattached hidden sessions in the same group
  85. if BYOBU_BACKEND == "tmux":
  86. output = subprocess.Popen(["tmux", "list-sessions"], stdout=subprocess.PIPE).communicate()[0]
  87. if sys.stdout.encoding is None:
  88. output = output.decode("UTF-8")
  89. else:
  90. output = output.decode(sys.stdout.encoding)
  91. if not output:
  92. return
  93. # Find the master session to extract the group name. We use
  94. # the group number to be extra sure the right session is getting
  95. # killed. We don't want to accidentally kill the wrong one
  96. pattern = "^%s:.+\\((group [^\\)]+)\\).*$" % session_name
  97. master = re.search(pattern, output, re.MULTILINE)
  98. if not master:
  99. return
  100. # Kill all the matching hidden & unattached sessions
  101. pattern = "^_%s-\\d+:.+\\(%s\\)$" % (session_name, master.group(1))
  102. for s in re.findall(pattern, output, re.MULTILINE):
  103. subprocess.Popen(["tmux", "kill-session", "-t", s.split(":")[0]])
  104. def update_environment(session):
  105. backend, session_name = session.split("____", 2)
  106. for var in BYOBU_UPDATE_ENVVARS:
  107. value = os.getenv(var)
  108. if value:
  109. if backend == "tmux":
  110. cmd = ["tmux", "setenv", "-t", session_name, var, value]
  111. else:
  112. cmd = ["screen", "-S", session_name, "-X", "setenv", var, value]
  113. subprocess.call(cmd, stdout=open(os.devnull, "w"))
  114. def attach_session(session):
  115. update_environment(session)
  116. backend, session_name = session.split("____", 2)
  117. cull_zombies(session_name)
  118. # must use the binary, not the wrapper!
  119. if backend == "tmux":
  120. if reuse_sessions:
  121. os.execvp("tmux", ["tmux", "-u", "new-session", "-t", session_name, ";", "set-option", "destroy-unattached"])
  122. else:
  123. os.execvp("tmux", ["tmux", "-u", "attach", "-t", session_name])
  124. else:
  125. os.execvp("screen", ["screen", "-AOxRR", session_name])
  126. sessions = get_sessions()
  127. show_shell = os.path.exists("%s/.always-select" % (BYOBU_CONFIG_DIR))
  128. if len(sessions) > 1 or show_shell:
  129. sessions.append("NEW")
  130. text.append("Create a new Byobu session (%s)" % BYOBU_BACKEND)
  131. sessions.append("SHELL")
  132. text.append("Run a shell without Byobu (%s)" % SHELL)
  133. if len(sessions) > 1:
  134. sys.stdout.write("\nByobu sessions...\n\n")
  135. tries = 0
  136. while tries < 3:
  137. i = 1
  138. for s in text:
  139. sys.stdout.write(" %d. %s\n" % (i, s))
  140. i += 1
  141. try:
  142. try:
  143. user_input = input("\nChoose 1-%d [1]: " % (i - 1))
  144. except Exception:
  145. user_input = ""
  146. if not user_input or user_input == "":
  147. choice = 1
  148. break
  149. try:
  150. choice = int(user_input)
  151. except Exception:
  152. choice = int(eval(user_input))
  153. if choice >= 1 and choice < i:
  154. break
  155. else:
  156. tries += 1
  157. choice = -1
  158. sys.stderr.write("\nERROR: Invalid input\n")
  159. except KeyboardInterrupt:
  160. sys.stdout.write("\n")
  161. sys.exit(0)
  162. except Exception:
  163. if choice == "" or choice == -1:
  164. choice = 1
  165. break
  166. tries += 1
  167. choice = -1
  168. sys.stderr.write("\nERROR: Invalid input\n")
  169. elif len(sessions) == 1:
  170. # Auto-select the only session
  171. choice = 1
  172. if choice >= 1:
  173. if sessions[choice - 1] == "NEW":
  174. # Create a new session
  175. if BYOBU_BACKEND == "tmux":
  176. os.execvp("byobu", ["byobu", "new-session", SHELL])
  177. else:
  178. os.execvp("byobu", ["byobu", SHELL])
  179. elif sessions[choice - 1] == "SHELL":
  180. os.execvp(SHELL, [SHELL])
  181. else:
  182. # Attach to the chosen session; must use the binary, not the wrapper!
  183. attach_session(sessions[choice - 1])
  184. # No valid selection, default to the youngest session, create if necessary
  185. if BYOBU_BACKEND == "tmux":
  186. os.execvp("tmux", ["tmux"])
  187. else:
  188. os.execvp("screen", ["screen", "-AOxRR"])