terminal reader with pty spawn

terminal_reader.py

# Correct and working version

import argparse
import json
import os
import pty
import sys


class ScriptRecorder:
    def __init__(self, script_file):
        self.script = script_file
        self.in_command = True
        self.buffer = b''
        self.PROMPT_MARKER = b'%c'
        self.current_command = b''
        self.current_output = b''

    def write_json_block(self):
        if self.current_command:
            block = {
                "command": self.current_command.decode().strip(),
                "output": self.current_output.decode().strip()
            }
            self.script.write(json.dumps(block, indent=2).encode() + b'\n')
            self.script.flush()
            self.current_command = b''
            self.current_output = b''

    def read(self, fd):
        data = os.read(fd, 512)

        if self.PROMPT_MARKER in data:
            if not self.in_command:
                self.write_json_block()
            self.in_command = True
        elif self.in_command and (b'\r\n' in data or b'\n' in data):
            self.in_command = False

        # Remove our marker from the output
        data = data.replace(self.PROMPT_MARKER, b'')

        # Store command or output
        if self.in_command:
            self.current_command += data
        else:
            self.current_output += data

        return data

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-a', dest='append', action='store_true')
    parser.add_argument('-p', dest='use_python', action='store_true')
    parser.add_argument("-i", dest="print", action="store_true")
    parser.add_argument('filename', nargs='?', default='typescript')
    options = parser.parse_args()

    shell = sys.executable if options.use_python else os.environ.get('SHELL', 'sh')
    filename = options.filename
    mode = 'ab' if options.append else 'wb'

    # Modify PS1 to include our marker
    user_prompt_command = os.environ.get('CLA_USER_SHELL_PROMPT_COMMAND', r'[\u@\h \W]\$ ')
    os.environ['PROMPT_COMMAND'] = f'{user_prompt_command}%c'

    with open(filename, mode) as script:
        recorder = ScriptRecorder(script)

        pty.spawn(shell, recorder.read)

        # Write final block if exists
        recorder.write_json_block()

if __name__ == '__main__':
    main()