Choggbuster/dvdread.py

import cffi
import os
import functools
import binascii
from datetime import timedelta


def loadlib(dll_path, *includes, **kwargs):
    ffi = cffi.FFI()
    for include in includes:
        ffi.cdef(open(include).read(), **kwargs)
    return ffi, ffi.dlopen(dll_path)


class DVDRead(object):
    def __init__(self, path, verbose="0", method="disc"):
        if verbose is None:
            os.environ.pop("DVDCSS_VERBOSE", None)
        else:
            os.environ["DVDCSS_VERBOSE"] = str(verbose)
        os.environ["DVDCSS_METHOD"] = method
        self.dvd = None
        self.ffi, self.lib = loadlib(
            "libdvdread-8.dll",
            "dvd_types.h",
            "dvd_reader.h",
            "ifo_types.h",
            "ifo_read.h",
            "ifo_print.h",
            "nav_types.h",
            "nav_read.h",
            "nav_print.h",
            packed=True,
        )
        self.path = path
        self.titles = {}
        self.open(path)
        self.lb_len=self.lib.DVD_VIDEO_LB_LEN

    def grab_ifo(self,title,bup=False):
        from tqdm import tqdm
        buf = self.ffi.new("unsigned char[]", 512)
        if bup:
            fh = self.lib.DVDOpenFile(self.dvd, title, self.lib.DVD_READ_INFO_BACKUP_FILE)
        else:
            fh = self.lib.DVDOpenFile(self.dvd, title, self.lib.DVD_READ_INFO_FILE)
        total_size = self.lib.DVDFileSize(fh)*self.lb_len
        remaining = total_size
        num_read = True
        pbar = tqdm(total=total_size, unit="iB", unit_scale=True, unit_divisor=1024,leave=False)
        while num_read:
            num_read=self.lib.DVDReadBytes( fh, buf, 512)
            num_read=min(num_read,remaining)
            remaining-=num_read
            pbar.update(num_read)
            yield self.ffi.buffer(buf,num_read)[:]
        self.lib.DVDCloseFile(fh)

    def grab_ifos(self):
        vmg_ifo = self.lib.ifoOpen(self.dvd, 0)
        if vmg_ifo == self.ffi.NULL:
            return
        title_sets = vmg_ifo.vts_atrt.nr_of_vtss
        for t in range(1,title_sets + 1):
            vts = self.lib.ifoOpen(self.dvd, t)
            if vts == self.ffi.NULL:
                continue
            self.lib.ifoClose(vts)
            outfile=os.path.join("RIP",f"VTS_{t:02}_0.ifo")
            with open(outfile, "wb") as out_ifo:
                for block in self.grab_ifo(t,bup=False):
                    out_ifo.write(block)
            outfile=os.path.join("RIP",f"VTS_{t:02}_0.bup")
            with open(outfile, "wb") as out_ifo:
                for block in self.grab_ifo(t,bup=True):
                    out_ifo.write(block)
        self.lib.ifoClose(vmg_ifo)

    def grab_vob(self,title):
        from tqdm import tqdm
        buf = self.ffi.new("unsigned char[]", 512 * self.lb_len)
        fh = self.lib.DVDOpenFile(self.dvd, title, self.lib.DVD_READ_TITLE_VOBS)
        total_size = self.lib.DVDFileSize(fh)*self.lb_len
        remaining = total_size
        num_read = True
        pos=0
        pbar = tqdm(total=total_size, unit="iB", unit_scale=True, unit_divisor=1024,leave=False)
        while remaining:
            num_read=self.lib.DVDReadBlocks( fh,pos, 512, buf)
            if num_read<0:
                raise RuntimeError("Error reading!")
            num_read_bytes=num_read*self.lb_len
            num_read_bytes=min(num_read_bytes,remaining)
            remaining-=num_read_bytes
            pbar.update(num_read_bytes)
            yield self.ffi.buffer(buf,num_read_bytes)[:]
            pos+=num_read
        pbar.close()
        self.lib.DVDCloseFile(fh)

    def grab_vobs(self):
        vmg_ifo = self.lib.ifoOpen(self.dvd, 0)
        if vmg_ifo == self.ffi.NULL:
            return
        title_sets = vmg_ifo.vts_atrt.nr_of_vtss
        for t in range(1,title_sets + 1):
            vts = self.lib.ifoOpen(self.dvd, t)
            if vts == self.ffi.NULL:
                continue
            self.lib.ifoClose(vts)
            outfile=os.path.join("RIP",f"VTS_{t:02}_0.vob")
            with open(outfile, "wb") as out_ifo:
                for block in self.grab_vob(t):
                    out_ifo.write(block)
        self.lib.ifoClose(vmg_ifo)


    def test(self):
        from tqdm import tqdm
        fn = 0
        chunk_size = 2048
        buf = self.ffi.new("unsigned char[]", chunk_size * 2048)
        for fn in range(title_sets + 1):
            pos = 0
            fh = self.lib.DVDOpenFile(self.dvd, fn, self.lib.DVD_READ_TITLE_VOBS)
            if fh:
                total_size = self.lib.DVDFileSize(fh)
                if total_size == -1:
                    self.lib.DVDCloseFile(fh)
                    break
                pbar = tqdm(total=total_size * 2048, unit="iB", unit_scale=True, unit_divisor=1024,leave=False)
                last=False
                with open(f"out_{fn}.vob", "wb") as out_vob:
                    while True:
                        if (pos+chunk_size)>total_size:
                            chunk_size=total_size-pos
                        count = self.lib.DVDReadBlocks(fh, pos, chunk_size, buf)
                        if count == -1:
                            break
                        pbar.update(
                            out_vob.write(self.ffi.buffer(buf, count * 2048)[:])
                        )
                        pos += count
                        if pos>=total_size:
                            break
                self.lib.DVDCloseFile(fh)
            fn += 1
            if fn>200:
                break

    def __del__(self):
        if self.dvd:
            self.lib.DVDClose(self.dvd)

    def open(self, path):
        # self.dvd_css=self.css_lib.dvdcss_open()
        self.dvd = self.lib.DVDOpen(bytes(path, "utf8"))
        vol_id = self.ffi.new("unsigned char[]", 32)
        self.lib.DVDDiscID(self.dvd, vol_id)
        self.disc_id = str(binascii.hexlify(self.ffi.buffer(vol_id, 16)[:]), "utf8")
        self.lib.DVDUDFVolumeInfo(self.dvd, vol_id, 32, self.ffi.NULL, 0)
        self.udf_disc_name = str(self.ffi.string(vol_id), "utf8")
        self.lib.DVDISOVolumeInfo(self.dvd, vol_id, 32, self.ffi.NULL, 0)
        self.iso_disc_name = str(self.ffi.string(vol_id), "utf8")
        self.ffi.release(vol_id)