198 lines
4.3 KiB
Python
198 lines
4.3 KiB
Python
import subprocess
|
|
import json
|
|
import sys
|
|
import os
|
|
import time
|
|
import tempfile
|
|
import base64
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
Recording_flag = "/tmp/screenshot_active"
|
|
Recordings = Path.home() / "records"
|
|
Tmp = tempfile.gettempdir()
|
|
|
|
def rand12():
|
|
return base64.urlsafe_b64encode(os.urandom(9)).decode()[:12]
|
|
|
|
|
|
def run(cmd):
|
|
return subprocess.run(cmd)
|
|
|
|
def find_focused_pid(node):
|
|
if isinstance(node, dict):
|
|
if node.get("focused") is True:
|
|
return node.get("pid")
|
|
|
|
for key in ("nodes", "floating_nodes"):
|
|
for child in node.get(key, []):
|
|
pid = find_focused_pid(child)
|
|
if pid:
|
|
return pid
|
|
|
|
elif isinstance(node, list):
|
|
for item in node:
|
|
pid = find_focused_pid(item)
|
|
if pid:
|
|
return pid
|
|
|
|
return None
|
|
|
|
def force_kill():
|
|
tree = json.loads(subprocess.check_output(
|
|
["swaymsg", "-t", "get_tree"],
|
|
text=True
|
|
))
|
|
pid = find_focused_pid(tree)
|
|
|
|
if not pid:
|
|
print("No focused window PID found")
|
|
return
|
|
|
|
os.kill(pid, 9)
|
|
|
|
def is_recording():
|
|
return subprocess.run(
|
|
["pgrep", "-x", "wf-recorder"],
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL
|
|
).returncode == 0
|
|
|
|
|
|
def record(fullscreen=False):
|
|
Recordings.mkdir(parents=True, exist_ok=True)
|
|
|
|
if is_recording():
|
|
subprocess.run(["killall", "wf-recorder"])
|
|
try:
|
|
os.remove(Recording_flag)
|
|
except FileNotFoundError:
|
|
pass
|
|
return
|
|
|
|
Path(Recording_flag).write_text("1")
|
|
|
|
slurp_cmd = ["slurp"]
|
|
if fullscreen:
|
|
slurp_cmd.append("-o")
|
|
|
|
selection = subprocess.getoutput(" ".join(slurp_cmd))
|
|
|
|
if not selection.strip():
|
|
return
|
|
|
|
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
|
output = Recordings / f"{timestamp}.mp4"
|
|
|
|
proc = subprocess.Popen(
|
|
["wf-recorder", "-g", selection, "-f", str(output)]
|
|
)
|
|
proc.wait()
|
|
|
|
if not output.exists():
|
|
return
|
|
|
|
compressed = output.with_name(output.stem + "_comp720p.mp4")
|
|
|
|
subprocess.run([
|
|
"ffmpeg",
|
|
"-i", str(output),
|
|
"-vf", "scale=-2:720",
|
|
"-vcodec", "libx264",
|
|
"-crf", "32",
|
|
"-preset", "fast",
|
|
str(compressed)
|
|
])
|
|
|
|
try:
|
|
output.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
def resolve_output(out):
|
|
if not out:
|
|
out = Tmp
|
|
|
|
out = os.path.expanduser(out)
|
|
|
|
if os.path.isdir(out) or out.endswith("/"):
|
|
os.makedirs(out, exist_ok=True)
|
|
return os.path.join(out, f"{rand12()}.png")
|
|
|
|
parent = os.path.dirname(out)
|
|
if parent:
|
|
os.makedirs(parent, exist_ok=True)
|
|
|
|
if not out.endswith(".png"):
|
|
out += ".png"
|
|
|
|
return out
|
|
|
|
|
|
def take_screenshot(output_path):
|
|
Path(Recording_flag).write_text("1")
|
|
|
|
try:
|
|
time.sleep(0.2)
|
|
|
|
freeze = subprocess.Popen(["wayfreeze"])
|
|
time.sleep(0.1)
|
|
|
|
try:
|
|
region = subprocess.check_output(["slurp"]).decode().strip()
|
|
|
|
subprocess.run(
|
|
["grim", "-g", region, output_path],
|
|
check=True
|
|
)
|
|
|
|
except subprocess.CalledProcessError:
|
|
return False
|
|
|
|
finally:
|
|
freeze.terminate()
|
|
|
|
if os.path.exists(output_path):
|
|
with open(output_path, "rb") as f:
|
|
subprocess.run(
|
|
["wl-copy", "--type", "image/png"],
|
|
input=f.read()
|
|
)
|
|
|
|
subprocess.run(["notify-send", "screenshot", f"Saved: {output_path}"])
|
|
return True
|
|
|
|
finally:
|
|
try:
|
|
os.remove(Recording_flag)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
return False
|
|
|
|
def main():
|
|
if len(sys.argv) < 2:
|
|
sys.exit(1)
|
|
|
|
cmd = sys.argv[1]
|
|
|
|
if cmd == "kill":
|
|
force_kill()
|
|
|
|
elif cmd == "record":
|
|
fullscreen = "--fullscreen" in sys.argv
|
|
record(fullscreen)
|
|
|
|
elif cmd == "screenshot":
|
|
out = None
|
|
|
|
if "--out" in sys.argv:
|
|
idx = sys.argv.index("--out")
|
|
if idx + 1 < len(sys.argv):
|
|
out = sys.argv[idx + 1]
|
|
|
|
output = resolve_output(out)
|
|
take_screenshot(output)
|
|
|
|
if __name__ == "__main__":
|
|
main() |