Files

193 lines
4.2 KiB
Python

import subprocess
import json
import sys
import os
import time
import tempfile
import base64
from datetime import datetime
from pathlib import Path
from warehouse import Media_flag, rand12
Recordings = Path.home() / "records"
Tmp = tempfile.gettempdir()
def run(cmd):
return subprocess.Popen(c, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
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:
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(Media_flag)
except FileNotFoundError:
pass
return
Path(Media_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(Media_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(Media_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()