fastapi_dvrip/nvr_types.py

from datetime import datetime
from asyncio_dvrip import DVRIPCam
import json
import struct
import base64

NVR_DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"

PRIMARY_STREAM = 0
SECONDARY_STREAM = 1
UNKNOWN_STREAM = 2

H264 = "h264"
H265 = "h265"
H265X = "h265x"

class File:
    def __init__(self, data, channel = 0, stream = 0) -> None:
        self.begin = datetime.strptime(data.get("BeginTime", data.get("begin")), NVR_DATETIME_FORMAT)
        self.end = datetime.strptime(data.get("EndTime", data.get("end")), NVR_DATETIME_FORMAT)
        self.DiskNo = data.get("DiskNo")
        self.SerialNo = data.get("SerialNo")
        self.size = int(data.get("FileLength"), 0) * 1024 if "FileLength" in data else data.get("size", 0)
        self.filename = data.get("FileName", data.get("filename"))
        self.filename_cleared = self.filename.split("/")[-1].replace("[", "_").replace("]", "_").replace("@","_")
        self.channel = channel
        self.stream = stream

    def __str__(self) -> str:
        return f"{self.filename_cleared}/{self.size}/{self.channel}/{self.stream}"

    def __repr__(self) -> str:
        return self.__str__()

    @property
    def to_b64(self):
        dict_obj = dict(self.__dict__)
        dict_obj["begin"] = self.begin.strftime(NVR_DATETIME_FORMAT)
        dict_obj["end"] = self.end.strftime(NVR_DATETIME_FORMAT)
        return base64.b64encode(json.dumps(dict_obj).encode()).decode("utf8")

    @property
    def json(self):
        return {"filename": self.filename_cleared, "size": self.size, "b64": self.to_b64.replace("==", "")}

    @staticmethod
    def from_b64(b64):
        data = json.loads(base64.b64decode(b64).decode('utf-8'))
        print(data)
        return File(data, data.get("channel"), data.get("stream"))

    async def generate_first_bytes(self, client:DVRIPCam, version = 0):
        client.logger.debug("init request")
        #init request
        await client.send(
            1424,
            {
                "Name": "OPPlayBack",
                "OPPlayBack": {
                    "Action": "Claim",
                    "Parameter": {
                        "PlayMode": "ByName",
                        "FileName": self.filename,
                        "StreamType": self.stream,
                        "Value": 0,
                        "TransMode": "TCP",
                    },
                    "StartTime": self.begin.strftime(NVR_DATETIME_FORMAT),
                    "EndTime": self.end.strftime(NVR_DATETIME_FORMAT),
                },
            },
        )
        client.logger.debug("download request")
        #download request
        msg = 1420
        data = {
                "Name": "OPPlayBack",
                "OPPlayBack": {
                    "Action": "DownloadStart",
                    "Parameter": {
                        "PlayMode": "ByName",
                        "FileName": self.filename,
                        "StreamType": self.stream,
                        "Value": 0,
                        "TransMode": "TCP",
                    },
                    "StartTime": self.begin.strftime(NVR_DATETIME_FORMAT),
                    "EndTime": self.end.strftime(NVR_DATETIME_FORMAT),
                },
            }

        #if client.socket_writer is None:
        #    client.logger.debug("socket writer is null")
        #    await client.connect()

        client.logger.debug("Blocking busy")
        await client.busy.acquire()
        if hasattr(data, "__iter__"):
            if version == 1:
                data["SessionID"] = f"{client.session:#0{12}x}"
            data = bytes(
                json.dumps(data, ensure_ascii=False, separators=(",", ":")), "utf-8"
            )

        tail = b"\x00"
        if version == 0:
            tail = b"\x0a" + tail
        pkt = (
            struct.pack(
                "BB2xII2xHI",
                255,
                version,
                client.session,
                client.packet_count,
                msg,
                len(data) + len(tail),
            )
            + data
            + tail
        )
        client.logger.debug("Send first package")
        client.socket_send(pkt)
        client.logger.debug("Grab first response")
        data = await client.socket_recv(20)
        client.logger.debug(data)
        if data is None or len(data) < 20:
            return None
        (
            head,
            version,
            client.session,
            sequence_number,
            msgid,
            len_data,
        ) = struct.unpack("BB2xII2xHI", data)
        return len_data

    async def get_file_stream(self, client: DVRIPCam, first_chunk_size) -> bytes:
        yield bytes(await client.receive_with_timeout(first_chunk_size))
        while True:
            header = await client.receive_with_timeout(20)
            len_data = struct.unpack("I", header[16:])[0]
            
            if len_data == 0:
                break

            yield bytes(await client.receive_with_timeout(len_data))
        await client.busy.release()
        await client.send(
                1420,
                {
                    "Name": "OPPlayBack",
                    "OPPlayBack": {
                        "Action": "DownloadStop",
                        "Parameter": {
                            "FileName": self.filename,
                            "PlayMode": "ByName",
                            "StreamType": self.stream,
                            "TransMode": "TCP",
                            "Channel": self.channel,
                            "Value": 0,
                        },
                        "StartTime": self.begin.strftime(NVR_DATETIME_FORMAT),
                        "EndTime": self.end.strftime(NVR_DATETIME_FORMAT),
                    },
                },
            )
        yield bytes(b"")

async def list_local_files(client: DVRIPCam, startTime, endTime, filetype, channel = 0, streamType = 0):
    # 1440 OPFileQuery
    result = []
    data = await client.send(
        1440,
        {
            "Name": "OPFileQuery",
            "OPFileQuery": {
                "BeginTime": startTime,
                "Channel": channel,
                "DriverTypeMask": "0x0000FFFF",
                "EndTime": endTime,
                "Event": "*",
                "StreamType": f"0x0000000{streamType}",
                "Type": filetype,
            },
        },
    )

    if data == None:
        client.logger.debug("Could not get files.")
        raise ConnectionRefusedError("Could not get files")

    # When no file can be found
    if data["Ret"] != 100:
        client.logger.debug(
            f"No files found for channel {channel} for this time range. Start: {startTime}, End: {endTime}"
        )
        return []

    # OPFileQuery only returns the first 64 items
    # we therefore need to add the results to a list, modify the starttime with the begintime value of the last item we received and query again
    # max number of results are 511
    result = data["OPFileQuery"]

    max_event = {"status": "init", "last_num_results": 0}
    while max_event["status"] == "init" or max_event["status"] == "limit":
        if max_event["status"] == "init":
            max_event["status"] = "run"
        while len(data["OPFileQuery"]) == 64 or max_event["status"] == "limit":
            newStartTime = data["OPFileQuery"][-1]["BeginTime"]
            data = await client.send(
                1440,
                {
                    "Name": "OPFileQuery",
                    "OPFileQuery": {
                        "BeginTime": newStartTime,
                        "Channel": channel,
                        "DriverTypeMask": "0x0000FFFF",
                        "EndTime": endTime,
                        "Event": "*",
                        "StreamType": "0x00000000",
                        "Type": filetype,
                    },
                },
            )
            result += data["OPFileQuery"]
            max_event["status"] = "run"

        if len(result) % 511 == 0 or max_event["status"] == "limit":
            client.logger.debug("Max number of events reached...")
            if len(result) == max_event["last_num_results"]:
                client.logger.debug(
                    "No new events since last run. All events queried"
                )
                return result

            max_event["status"] = "limit"
            max_event["last_num_results"] = len(result)

    client.logger.debug(f"Found {len(result)} files.")
    return result