DRB-CnC/Client/pdab/recorder.py

import pyaudio
import wave, logging, threading, time, queue, signal, argparse, audioop
from os import path, makedirs

logging.basicConfig(format="%(asctime)s: %(message)s", level=logging.INFO,datefmt="%H:%M:%S")

class DiscordRecorder:
    def __init__(self, DEVICE_ID, CHUNK = 960, FORMAT = pyaudio.paInt16, CHANNELS = 2, RATE = 48000, FILENAME = "./recs/radio.wav"):
        self.pa_instance = pyaudio.PyAudio()

        self.DEVICE_ID = DEVICE_ID
        self.CHUNK = CHUNK
        self.FORMAT = FORMAT
        self.CHANNELS = CHANNELS
        self.RATE = RATE
        self.NG_fadeout = 240/20    # Fadeout value used to hold the noisegate after de-triggering
        self.NG_fadeout_count = 0   # A count set when the noisegate is triggered and was de-triggered
        self.process_set_count = 0  # Counts how many processes have been made

        self.FILENAME = FILENAME
        self._check_file_path_exists()

        self.queued_frames = queue.Queue()

        self.stop_threads = threading.Event()

        self.recording_thread = None
        self.saving_thread = None

        self.running_stream = None

    # Wrapper to check if the given filepath (not file itself) exists
    def _check_file_path_exists(self):
        if not path.exists(path.dirname(self.FILENAME)):            
            makedirs(path.dirname(self.FILENAME), exist_ok=True)

    # Wrapper for the recorder thread; Adds new data to the queue
    def _recorder(self):
        logging.info("* Recording Thread Starting")
        while True:
            try:
                curr_buffer = bytearray(self.stream.stream.read(self.CHUNK))
                buffer_rms = audioop.rms(curr_buffer, 2)
                if buffer_rms > 0:
                    buffer_decibel = 20 * math.log10(buffer_rms)

                    if self.process_set_count % 10 == 0:
                        if buffer_decibel >= self.stream.THRESHOLD:
                            LOGGER.debug(f"[Noisegate Open] {buffer_decibel} db")
                        else:
                            LOGGER.debug(f"[Noisegate Closed] {buffer_decibel} db")

                    if buffer_decibel >= self.stream.THRESHOLD:
                        self.NG_fadeout_count = self.NG_fadeout
                        self.process_set_count += 1
                        if curr_buffer:
                            return self.queued_frames.put(curr_buffer)

                    else:
                        if self.NG_fadeout_count > 0:
                            self.NG_fadeout_count -= 1
                            LOGGER.debug(f"Frames in fadeout remaining: {self.NG_fadeout_count}")
                            self.process_set_count += 1
                            if curr_buffer:
                                return self.queued_frames.put(curr_buffer)
            
            except OSError as e:
                LOGGER.warning(e)
                pass

            # check for stop
            if self.stop_threads.is_set():
                break

    # Wrapper for saver thread; Saves the queue to the file
    def _saver(self):
        logging.info("* Saving Thread Starting")
        while True:
            if not self.queued_frames.empty():
                dequeued_frames = []
                for i in range(self.queued_frames.qsize()):
                    dequeued_frames.append(self.queued_frames.get())

                if not path.isfile(self.FILENAME):
                    wf = wave.open(self.FILENAME, 'wb')
                    wf.setnchannels(self.CHANNELS)
                    wf.setsampwidth(self.pa_instance.get_sample_size(self.FORMAT))
                    wf.setframerate(self.RATE)
                    wf.writeframes(b''.join(dequeued_frames))
                    wf.close()
                else:                                
                    read_file = wave.open(self.FILENAME, 'rb')
                    read_file_data = read_file.readframes(read_file.getnframes())
                    read_file.close()
                        
                    wf = wave.open(self.FILENAME, 'wb')
                    wf.setnchannels(self.CHANNELS)
                    wf.setsampwidth(self.pa_instance.get_sample_size(self.FORMAT))
                    wf.setframerate(self.RATE)

                    wf.writeframes(read_file_data)
                    wf.writeframes(b''.join(dequeued_frames))
                    wf.close()

            # check for stop
            if self.stop_threads.is_set():
                break

            time.sleep(5)

    # Start the recording function
    def start_recording(self):
        logging.info("* Recording")

        self.running_stream = self.pa_instance.open(
            input_device_index=self.DEVICE_ID,
            format=self.FORMAT,
            channels=self.CHANNELS,
            rate=self.RATE,
            input=True,
            frames_per_buffer=self.CHUNK
        )
        
        self.recording_thread = threading.Thread(target=self._recorder)
        self.recording_thread.start()

        self.saving_thread = threading.Thread(target=self._saver)
        self.saving_thread.start()
    
    # Stop the recording function
    def stop_recording(self):
        self.stop_threads.set()
        self.recording_thread.join()
        self.saving_thread.join()
        self.running_stream.stop_stream()
        self.running_stream.close()
        self.pa_instance.terminate()
        
        logging.info("* Done recording")


class GracefulExitCatcher:
  def __init__(self, stop_callback):
    self.stop = False

    # The function to run when the exit signal is caught
    self.stop_callback = stop_callback

    # Update what happens when these signals are caught
    signal.signal(signal.SIGINT, self.exit_gracefully)
    signal.signal(signal.SIGTERM, self.exit_gracefully)

  def exit_gracefully(self, *args):
    logging.info("* Stop signal caught...")
    
    # Stop the main loop
    self.stop = True

    # Run the given callback function
    self.stop_callback()




if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("deviceId", type=int, help="The ID of the audio device to use")
    parser.add_argument("filename", type=str, help="The filepath/filename of the output file")    
    args = parser.parse_args()

    logging.debug("Arguments:", args)
    
    recorder = DiscordRecorder(args.deviceId, FILENAME=args.filename)

    exit_catcher = GracefulExitCatcher(recorder.stop_recording)

    recorder.start_recording()

    while not exit_catcher.stop:
        time.sleep(1)