import datetime
import os

from bitstring import ConstBitStream


class HikvisionDrive:

    def __init__(self, file_loc):
        # create data stream of hard drive
        self.file = file_loc
        self.data = ConstBitStream(filename=self.file)

        # get information from master sector
        self.data.pos = 584 * 8
        self.capacity = self.data.read('uintle:64')

        self.data.pos += 56 * 8
        self.data_block_size = self.data.read('uintle:64')
        self.total_data_blocks = self.data.read('uintle:32')

        self.data.pos += 4 * 8
        self.hikbtree_offset = self.data.read('uintle:64')
        self.hikbtree_size = self.data.read('uintle:32')

        # get information from hikbtree
        self.data.pos = (self.hikbtree_offset + 80) * 8
        self.page_list_offset = self.data.read('uintle:64')

        # initialise list for data blocks
        self.data_blocks = []

    def show_info(self):
        print('-Hard Drive Information-')
        print('File: ' + str(self.file))
        print('Capacity: ' + str(self.capacity) + ' bytes')
        print('Data Block Size: ' + str(self.data_block_size) + ' bytes')
        print('Number of Data Blocks: ' + str(self.total_data_blocks))
        print('HikBTree Offset: ' + str(self.hikbtree_offset))
        print('HikBTree Size: ' + str(self.hikbtree_size) + ' bytes')
        print('Page List Offset: ' + str(self.page_list_offset))

    def save_blocks(self, save_directory):
        # loop through data blocks and save the data into a structured directory
        # setup progress counter
        x = 1
        with open(self.file, 'rb') as f1:
            for data_block in self.data_blocks[:1000]:
                # update progress
                print('saving data block #', x, 'of', self.total_data_blocks, '...')
                # set directory
                date = data_block['start_time']
                save_dir = save_directory + '\' + str(date.year) + '-' + date.strftime('%B') + '\'
                    + date.strftime('%d')
                if not os.path.exists(save_dir):
                    os.makedirs(save_dir)
                # save the data to a file
                file_name = 'CH' + str(data_block['channel']) + '_' + date.strftime('%H') + '-' + date.strftime(
                    '%M') + '-' + date.strftime('%S') + '_block' + str(x) + '.raw'
                length = self.data_block_size
                f1.seek(data_block['offset'])
                with open(str(save_dir + '\' + file_name), 'wb') as f2:
                    while length:
                        chunk = min(1024 * 1024, length)
                        data = f1.read(chunk)
                        f2.write(data)
                        length -= chunk
                f2.close()
                x += 1
            f1.close()

    def scan_pages(self):
        # create list of page offsets
        page_offsets = []
        # get the offset to Page #1
        self.data.pos = (self.page_list_offset + 24) * 8
        page_offsets.append(self.data.read('uintle:64'))
        # loop through the rest of the list
        self.data.pos += 64 * 8
        page_offset = self.data.read('uintle:64')
        while page_offset != 0:
            page_offsets.append(page_offset)
            self.data.pos += 40 * 8
            page_offset = self.data.read('uintle:64')
        # create list of data blocks from the pages
        for page in page_offsets:
            self.data.pos = (page + 96) * 8
            entry = True if self.data.read('hex:64') == 'ffffffffffffffff' else False
            while entry:
                video_exists = True if self.data.read('hex:64') == '0000000000000000' else False
                self.data.pos += 1 * 8
                channel = self.data.read('uint:8')
                self.data.pos += 6 * 8
                start_time = datetime.datetime.fromtimestamp(self.data.read('uintle:32'))
                end_time = datetime.datetime.fromtimestamp(self.data.read('uintle:32'))
                data_offset = self.data.read('uintle:64')
                self.data_blocks.append({'channel': channel, 'has_video': video_exists, 'start_time': start_time,
                                         'end_time': end_time, 'offset': data_offset})
                # move to next block entry
                self.data.pos += 8 * 8
                entry = True if self.data.read('hex:64') == 'ffffffffffffffff' else False


drive = HikvisionDrive(r'D:213_HDD.001')
drive.show_info()
drive.scan_pages()
drive.save_blocks(r'D:Carved Files')