find_duplicates.py

#!python3.8


# NOTE: This version is incomplete and will refuse to run properly. See previous version for stable, working code.


import csv
import mmap
import time
import shutil
import struct
import sqlite3
import threading
from pathlib import Path
from functools import partial
from datetime import datetime
from hashlib import blake2b as _blake2b
from dataclasses import dataclass, astuple


DATABASE_FILE = 'duplicates.db'
CSV_FILE = 'duplicates.csv'
HASH_SIZE = 16
HASH_FUNCTION = partial(_blake2b, digest_size=HASH_SIZE)
DEFAULT_HASH = HASH_FUNCTION().digest()  # used for empty files
READ_SIZE = 65536


@dataclass
class File():
    path          : Path
    last_modified : datetime = None
    size          : int      = None
    hash          : bytes    = None
    symlink_to    : str      = None
    shortcut_to   : str      = None

    def __post_init__(self):
        """Clean up the data and add any missing data"""
        path = Path(self.path).absolute()
        self.path = str(path)
        if path.is_symlink():
            self.symlink_to = str(path.resolve())
            return

        self.last_modified = self.last_modified or \
            datetime.fromtimestamp(path.stat().st_mtime)
        self.size = self.size or path.stat().st_size


    def compute_hash(self):
        """Hashshes are only computed if you want to them to be"""
        self.hash = hash_file(self.path)


    def as_tuple(self):
        return astuple(self)

    def is_symlink(self):
        return bool(self.symlink_to)

    def is_shortcut(self):
        return bool(self.shortcut_to)


class FileDB(sqlite3.Connection):
    def __init__(self):
        super().__init__(
            DATABASE_FILE,
            detect_types=sqlite3.PARSE_DECLTYPES,
            check_same_thread=False,
        )

        self.executescript("""
            CREATE TABLE IF NOT EXISTS files (
                path TEXT UNIQUE NOT NULL,
                last_modified TIMESTAMP,
                size INT,
                hash BLOB,
                symlink_to TEXT,
                shortcut_to TEXT
            );
            
            CREATE VIEW IF NOT EXISTS duplicates AS 
                SELECT 
                     path AS Path
                    ,last_modified AS LastModified
                    ,size AS FileSize
                    ,ttl_size AS TotalSize
                    ,HEX(hash) AS Hash
                    
                FROM (
                    SELECT
                         path 
                        ,last_modified
                        ,size
                        ,SUM(size) OVER(PARTITION BY hash) as ttl_size
                        ,hash
                    FROM (
                        SELECT
                             path
                            ,last_modified
                            ,size
                            ,hash
                            ,COUNT(hash) OVER (PARTITION BY hash) AS count_hash
                        FROM files
                        WHERE symlink_to is NULL and shortcut_to is NULL
                    )
                    WHERE count_hash > 1
                )
                ORDER BY ttl_size DESC
        ;
        """)
        self.commit()


    def get_file(self, path):
        return self.execute('select * from files where path=?', (path,)).fetchone()


    def insert_file(self, file):
        if isinstance(file, File):
            file = file.as_tuple()
        try:
            self.execute("""
                INSERT OR REPLACE INTO files(
                    path, last_modified, size, hash, symlink_to, shortcut_to
                ) VALUES (?,?,?,?,?,?)
            """, file)
        except:
            breakpoint()


    def check_duplicates(self):
        return bool(
            self.execute(
                'select * from duplicates limit 1'
            ).fetchone()
        )


    def get_duplicates(self):
        c = self.execute('select * from duplicates')
        headers = [i[0] for i in c.description]
        yield headers
        for row in c:
            yield row


def get_database():
    return FileDB()


def hash_file(path:str, file_size=None):
    """Compute the hash of a file using memory views
    optimized for the Windows platform. This function 
    can be run within its own thread.
    """
    if file_size is None:
        file_size = Path(path).stat().st_size


    # Empty files cannot be mmap'd on Windows
    if file_size == 0:
        return DEFAULT_HASH

    # reduce global lookups
    hash_function = HASH_FUNCTION
    ACCESS_READ = mmap.ACCESS_READ
    read_size = READ_SIZE

    hash = hash_function()
    offset = 0 
    with open(path, 'rb') as f:
        # if the file is too big to read all at once,
        # read it instead in READ_SIZE chunks using
        # a sliding window. Stop on the last window,
        # since it will be < READ_SIZE
        if file_size > read_size:
            for offset in range(0, file_size, read_size)[:-1]:
                with mmap.mmap(
                    f.fileno(), 
                    length=read_size, 
                    access=ACCESS_READ,
                    offset=offset,
                ) as m:   
                    hash.update(m)
            # set offset to start of next window.
            offset += read_size

        # get last window, OR, the whole file if no windowing was used.
        with mmap.mmap(
            f.fileno(), 
            length=(file_size - offset),
            access=mmap.ACCESS_READ,
            offset=offset,
        ) as m:   
            hash.update(m)

    return hash.digest()


def get_input(prompt):
    try:
        return input(prompt)
    except (KeyboardInterrupt, EOFError):
        return None


def b2h(n:int):
    """Converts bytes to human readable number"""
    if n < 1024:
        return f'{n} Bytes'
    factors = {
        'KiB': 1024,
        'MiB': 1024 ** 2,
        'GiB': 1024 ** 3,
        'TiB': 1024 ** 4,
        'PiB': 1024 ** 5,
    }
    for unit, factor in reversed(factors.items()):
        if n >= factor:
            return f'{n/factor:.2f} {unit}'


def decode_shortcut(path):
    """Given a file path to a shortcut (.lnk) file, 
    extract the destination path. See the documentation:
    https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-shllink/16cb4ca1-9339-4d0c-a68d-bf1d6cc0f943
    """
    
    raise NotImplementedError
    
    file_format_error = ValueError('Not a shortcut file')
    if not path.endswith('.lnk'):
        raise file_format_error
    with open(path, 'rb') as f:
        data = f.read()
    
    if not data.startswith(0x4C):  # ASCII 'L'
        raise file_format_error
    
    


def crawl_dir(path, progress_callback=None):

    raise NotImplementedError(
        'TODO: add support for shortcut files'
    )

    def print_progress():
        """Abuse race conditions to print the current path :)
        """
        nonlocal progress_callback
        nonlocal timer_lock
        nonlocal path
        nonlocal files_processed
        w, _ = shutil.get_terminal_size()
        p = str(path)
        if len(p) >= w-9:
            l = w//5
            p = p[:l] + '...' + p[l:]
        progress_callback(f'{files_processed} {p:<{w}}')
        timer_lock.release()

    db = get_database()
    files_processed = 0
    timer_lock = threading.Lock()
    timer = None

    for path in Path(path).glob('**/*'):
        if not path.is_file():
            continue

        file = File(path)
        if file.is_symlink():
            # do not compute hashes on symlinks
            db.insert_file(file.as_tuple())            

        # if it's a real file, it must be hashed
        else:
            if (f:=db.get_file(file.path)) is not None:
                db_file = File(*f)
                if file.last_modified != db_file.last_modified:
                    print('Changed detected:', file.path)
                    file.compute_hash()
                    db.insert_file(file.as_tuple())
            else:
                file.compute_hash()
                db.insert_file(file.as_tuple())

        files_processed += 1
        if files_processed % 1000 == 0:
            db.commit()

        if not timer_lock.locked():
            timer_lock.acquire()
            timer = threading.Timer(1, print_progress)
            timer.start()

    timer.cancel()
    db.commit()
    return files_processed


def check_duplicates():
    with get_database() as db:
        return db.check_duplicates()


def export_duplicates():
    with get_database() as db:
        dups = db.get_duplicates()
    with open(CSV_FILE, 'x') as f:
        w = csv.writer(f)
        for row in dups:
            w.writerow(row)


def cli():
    print('Enter a path to crawl. When done, use CTRL+C to check for duplicates.')
    while True:
        root_dir = get_input('Enter a directory> ')

        if root_dir is None:
            break  #  on empty input (user just pressed enter)

        if not root_dir:
            continue   # on ^C or ^Z

        if '..' in root_dir:
            print("Paths cannot be relative via ..")
            continue

        p = Path(root_dir)
        if p.exists():
            if p.is_dir():
                print('Crawling...')
                cb = partial(print, end='\r')
                files_processed = crawl_dir(root_dir, cb)
                print(f'\nProcessed {files_processed} files')
            else:
                print('Can only crawl directories')
        else:
            print('Directory not found')

    if check_duplicates():
        i = get_input('Duplicates found! Show them? y/[n] ')
        if i:
            if i.casefold().startswith('y'):
                while True:
                    print('Saving to', CSV_FILE)
                    try:
                        export_duplicates()
                        break
                    except FileExistsError:
                        i = get_input(
                            'ERROR: File exists! Rename or delete it. Press Enter '
                            'to try again, or CTRL + C to quit')
                        if i is None:
                            return
    print('Done.')


def main():
    cli()


if __name__ == '__main__':
    main()