Download rtmp stream file

rtmp_stream_downloader.py

import librtmp
import time
import re


class ParamsError(Exception):
    """Raised while download params are incorrect."""
    pass


class RtmpDownloader:
    """
    :param url: rtmp stream url
    :type url: `str`
    :param output_name: output file name, should have flv extension
    :type output_name: `str`
    """

    TIME_COUNTER_THRESH = 5
    RTMP_RE = r'^rtmp://[a-zA-Z0-9.:]+(/[a-z0-9]+){2}'
    NAME_RE = r'^([a-zA-Z0-9]+[-_]?)+(.flv){1}'

    def __init__(self, url, output_name):
        check_url = re.search(self.RTMP_RE, url)
        check_name = re.search(self.NAME_RE, output_name)

        if check_url is None:
            raise ParamsError("Stream url is not correct, it should be like rtmp://servername.com/app/filepath")

        if check_name is None:
            raise ParamsError("Name is incorrect, it should has .flv extension")

        self._url = url
        self._output_name = output_name
        self.stream = None

    def connect(self):
        """
        Establish a connection to rtmp server and create a stream
        """
        conn = librtmp.RTMP(self._url, live=True, timeout=1)
        try:
            # Attempt to connect
            conn.connect()
            print("INFO - Connected successfully to server")

            # Get a file-like object to access to teh stream
            self.stream = conn.create_stream()
            print("INFO - Created stream")
        except librtmp.RTMPError:
            print("ERROR - Cannot connect to stream")

    def start(self):
        self.connect()

        i = 1
        data = None
        is_recording = True
        timeout_counter = 0

        while True:
            # Read 1024 bytes of data
            data = self.stream.read(1024)
            print("INFO - Read chunk {}".format(i))
            i = i + 1

            if data == b'':
                time.sleep(1)
                timeout_counter = timeout_counter + 1
                self.connect()  # try to reconnect
            else:
                timeout_counter = 0

            if timeout_counter >= self.TIME_COUNTER_THRESH:
                break

            with open(self._output_name, "ab+") as f:
                if data != b'':
                    f.write(data)

        self.stream.close()