-
Notifications
You must be signed in to change notification settings - Fork 73
/
savescreen.py
executable file
·176 lines (154 loc) · 5.88 KB
/
savescreen.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
#!/usr/bin/env python3
#
# screen history saving script - saves your current windows and their titles,
# then appends this to ~/.screenrc and writes the result to ~/.screenrc.save
#
# This is intended to be run on a regular basis; I cron it every minute.
#
# WARNING - this expects only one screen session to be running as your user.
#
# Copyright 2014-2024 Jason Antman <[email protected]> <http://www.jasonantman.com>
# Use this however you want; send changes back to me please.
#
# The canonical version of this script is available at:
# <https://github.com/jantman/misc-scripts/blob/master/savescreen.py>
#
# CHANGELOG:
# 2014-07-24 Jason Antman <[email protected]>
# * first public version
#
# 2014-07-25 Jason Antman <[email protected]>
# * restored windows should go to their last PWD, set in ~/.bashrc
# (see <http://blog.jasonantman.com/2014/07/session-save-and-restore-with-bash-and-gnu-screen/>)
#
# 2015-03-26 Jason Antman <[email protected]>
# * add lockfile to prevent runaway cron processes
#
# 2024-01-23 Jason Antman <[email protected]>
# * completely rewrite for modern Python (>= 3.9)
# * remove ugly lockfile hackery
# * set 30-second hard timeout
####################################################
import subprocess
import os
from platform import node
from datetime import datetime
import sys
import logging
import argparse
import signal
import fcntl
from pathlib import Path
from typing import Optional
logging.basicConfig(
level=logging.WARNING,
format="[%(asctime)s %(levelname)s] %(message)s"
)
logger: logging.Logger = logging.getLogger()
class ScreenSaver:
LOCKFILE_PATH: str = os.path.abspath(os.path.expanduser('~/.savescreen.lock'))
def run(self, timeout: int = 30):
if not os.path.exists(self.LOCKFILE_PATH):
logger.info('Creating file: %s', self.LOCKFILE_PATH)
Path(self.LOCKFILE_PATH).touch()
def handler(*_):
raise TimeoutError()
# set the timeout handler
signal.signal(signal.SIGALRM, handler)
signal.alarm(timeout)
try:
ptr = os.open(self.LOCKFILE_PATH, os.O_WRONLY)
try:
fcntl.lockf(ptr, fcntl.LOCK_EX | fcntl.LOCK_NB)
except IOError:
logger.error(
'Unable to acquire lock on %s; already running?',
self.LOCKFILE_PATH, exc_info=True
)
return
self._do_it()
os.close(ptr)
except TimeoutError:
logger.error('Timeout at %d seconds', timeout)
finally:
signal.alarm(0)
def _get_windows(self, timeout: int = 25) -> Optional[str]:
cmd = ['screen', '-Q', 'windows', '%n %t|']
logger.debug('Running: %s', ' '.join(cmd))
p = subprocess.run(
cmd, stdout=subprocess.PIPE, timeout=timeout, text=True
)
logger.debug(
'Command exited %d with STDOUT: %s', p.returncode, p.stdout
)
if p.returncode != 0:
return None
return p.stdout
def _windowstr_to_dict(self, windowstr: str) -> dict:
"""This is really crap code, but it's been working for a decade..."""
windows = {}
for window in windowstr.split('|'):
if window.strip() == '':
continue
parts = window.strip().split(' ', maxsplit=1)
windows[int(parts[0])] = parts[1]
return windows
def _do_it(self):
windowstr: Optional[str] = self._get_windows()
if not windowstr:
return
windows: dict = self._windowstr_to_dict(windowstr)
logger.debug('Windows: %s', windows)
# read in screenrc
logger.debug("Reading .screenrc")
with open(os.path.expanduser('~/.screenrc'), 'r') as fh:
screenrc = fh.read()
# get rid of the first "local 0" line if it's there
screenrc = screenrc.replace("screen -t local 0\n", "")
logger.debug("Writing .screenrc.save")
# write it out to the save location, with the windows added
dirpath = os.path.expanduser('~/.screendirs')
with open(os.path.expanduser('~/.screenrc.save'), 'w') as fh:
fh.write(screenrc)
fh.write("\n\n")
dstr: str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
fh.write(f"# .screenrc.save generated on {node()} at {dstr}\n")
for n in range(0, max(windows.keys()) + 1):
fh.write(
f'screen -t "{windows.get(n, "bash")}" {n} sh -c "cd $(readlink -fn {dirpath}/{n}); bash"\n'
)
fh.write("\n")
def parse_args(argv):
p = argparse.ArgumentParser(description='Save screen session state to disk')
p.add_argument('-v', '--verbose', dest='verbose', action='store_true',
default=False, help='verbose output')
p.add_argument('-t', '--timeout', dest='maxtime', action='store',
default=30, help='Maximum runtime seconds (default: 30)')
args = p.parse_args(argv)
return args
def set_log_info(l: logging.Logger):
"""set logger level to INFO"""
set_log_level_format(
l,
logging.INFO,
'%(asctime)s %(levelname)s:%(name)s:%(message)s'
)
def set_log_debug(l: logging.Logger):
"""set logger level to DEBUG, and debug-level output format"""
set_log_level_format(
l,
logging.DEBUG,
"%(asctime)s [%(levelname)s %(filename)s:%(lineno)s - "
"%(name)s.%(funcName)s() ] %(message)s"
)
def set_log_level_format(lgr: logging.Logger, level: int, fmt: str):
"""Set logger level and format."""
formatter = logging.Formatter(fmt=fmt)
lgr.handlers[0].setFormatter(formatter)
lgr.setLevel(level)
if __name__ == "__main__":
args = parse_args(sys.argv[1:])
# set logging level
if args.verbose:
set_log_debug(logger)
ScreenSaver().run(timeout=args.maxtime)