CouchPotatoServer/couchpotato/core/plugins/renamer/main.py

from couchpotato import get_session
from couchpotato.api import addApiView
from couchpotato.core.event import addEvent, fireEvent, fireEventAsync
from couchpotato.core.helpers.encoding import toUnicode
from couchpotato.core.helpers.request import jsonified
from couchpotato.core.helpers.variable import getExt, mergeDicts
from couchpotato.core.logger import CPLog
from couchpotato.core.plugins.base import Plugin
from couchpotato.core.settings.model import Library, File
from couchpotato.environment import Env
import os
import re
import shutil
import traceback

log = CPLog(__name__)


class Renamer(Plugin):

    renaming_started = False

    def __init__(self):

        addApiView('renamer.scan', self.scanView, docs = {
            'desc': 'For the renamer to check for new files to rename',
        })

        addEvent('renamer.scan', self.scan)
        addEvent('app.load', self.scan)

        fireEvent('schedule.interval', 'renamer.scan', self.scan, minutes = self.conf('run_every'))

    def scanView(self):

        fireEventAsync('renamer.scan')

        return jsonified({
            'success': True
        })

    def scan(self):

        if self.isDisabled():
            return

        if self.renaming_started is True:
            log.error('Renamer is disabled to avoid infinite looping of the same error.')
            return

        # Check to see if the "to" folder is inside the "from" folder.
        if not os.path.isdir(self.conf('from')) or not os.path.isdir(self.conf('to')):
            log.debug('"To" and "From" have to exist.')
            return
        elif self.conf('from') in self.conf('to'):
            log.error('The "to" can\'t be inside of the "from" folder. You\'ll get an infinite loop.')
            return

        groups = fireEvent('scanner.scan', folder = self.conf('from'), single = True)

        self.renaming_started = True

        destination = self.conf('to')
        folder_name = self.conf('folder_name')
        file_name = self.conf('file_name')
        trailer_name = self.conf('trailer_name')
        nfo_name = self.conf('nfo_name')
        separator = self.conf('separator')

        db = get_session()

        for group_identifier in groups:

            group = groups[group_identifier]
            rename_files = {}
            remove_files = []
            remove_releases = []

            # Add _UNKNOWN_ if no library item is connected
            if not group['library']:
                if group['dirname']:
                    rename_files[group['parentdir']] = group['parentdir'].replace(group['dirname'], '_UNKNOWN_%s' % group['dirname'])
                else: # Add it to filename
                    for file_type in group['files']:
                        for rename_me in group['files'][file_type]:
                            filename = os.path.basename(rename_me)
                            rename_files[rename_me] = rename_me.replace(filename, '_UNKNOWN_%s' % filename)

            # Rename the files using the library data
            else:
                group['library'] = fireEvent('library.update', identifier = group['library']['identifier'], single = True)
                if not group['library']:
                    log.error('Could not rename, no library item to work with: %s' % group_identifier)
                    continue

                library = group['library']

                # Find subtitle for renaming
                fireEvent('renamer.before', group)

                # Remove weird chars from moviename
                movie_name = re.sub(r"[\x00\/\\:\*\?\"<>\|]", '', group['library']['titles'][0]['title'])

                # Put 'The' at the end
                name_the = movie_name
                if movie_name[:4].lower() == 'the ':
                    name_the = movie_name[4:] + ', The'

                replacements = {
                     'ext': 'mkv',
                     'namethe': name_the.strip(),
                     'thename': movie_name.strip(),
                     'year': library['year'],
                     'first': name_the[0].upper(),
                     'quality': group['meta_data']['quality']['label'],
                     'quality_type': group['meta_data']['quality_type'],
                     'video': group['meta_data'].get('video'),
                     'audio': group['meta_data'].get('audio'),
                     'group': group['meta_data']['group'],
                     'source': group['meta_data']['source'],
                     'resolution_width': group['meta_data'].get('resolution_width'),
                     'resolution_height': group['meta_data'].get('resolution_height'),
                }

                for file_type in group['files']:

                    # Move nfo depending on settings
                    if file_type is 'nfo' and not self.conf('rename_nfo'):
                        log.debug('Skipping, renaming of %s disabled' % file_type)
                        if self.conf('cleanup'):
                            for current_file in group['files'][file_type]:
                                remove_files.append(current_file)
                        continue

                    # Subtitle extra
                    if file_type is 'subtitle_extra':
                        continue

                    # Move other files
                    multiple = len(group['files']['movie']) > 1 and not group['is_dvd']
                    cd = 1 if multiple else 0

                    for current_file in sorted(list(group['files'][file_type])):

                        # Original filename
                        replacements['original'] = os.path.splitext(os.path.basename(current_file))[0]
                        replacements['original_folder'] = fireEvent('scanner.remove_cptag', group['dirname'], single = True)

                        # Extension
                        replacements['ext'] = getExt(current_file)

                        # cd #
                        replacements['cd'] = ' cd%d' % cd if cd else ''
                        replacements['cd_nr'] = cd

                        # Naming
                        final_folder_name = self.doReplace(folder_name, replacements)
                        final_file_name = self.doReplace(file_name, replacements)
                        replacements['filename'] = final_file_name[:-(len(getExt(final_file_name)) + 1)]

                        # Group filename without cd extension
                        replacements['cd'] = ''
                        replacements['cd_nr'] = ''
                        group['filename'] = self.doReplace(file_name, replacements)[:-(len(getExt(final_file_name)) + 1)]

                        # Meta naming
                        if file_type is 'trailer':
                            final_file_name = self.doReplace(trailer_name, replacements)
                        elif file_type is 'nfo':
                            final_file_name = self.doReplace(nfo_name, replacements)

                        # Seperator replace
                        if separator:
                            final_file_name = final_file_name.replace(' ', separator)

                        # Move DVD files (no structure renaming)
                        if group['is_dvd'] and file_type is 'movie':
                            found = False
                            for top_dir in ['video_ts', 'audio_ts', 'bdmv', 'certificate']:
                                has_string = current_file.lower().find(os.path.sep + top_dir + os.path.sep)
                                if has_string >= 0:
                                    structure_dir = current_file[has_string:].lstrip(os.path.sep)
                                    rename_files[current_file] = os.path.join(destination, final_folder_name, structure_dir)
                                    found = True
                                    break

                            if not found:
                                log.error('Could not determine dvd structure for: %s' % current_file)

                        # Do rename others
                        else:
                            if file_type is 'leftover':
                                if self.conf('move_leftover'):
                                    rename_files[current_file] = os.path.join(destination, final_folder_name, os.path.basename(current_file))
                            elif file_type not in ['subtitle']:
                                rename_files[current_file] = os.path.join(destination, final_folder_name, final_file_name)

                        # Check for extra subtitle files
                        if file_type is 'subtitle':

                            # rename subtitles with or without language
                            #rename_files[current_file] = os.path.join(destination, final_folder_name, final_file_name)
                            sub_langs = group['subtitle_language'].get(current_file, [])

                            rename_extras = self.getRenameExtras(
                                extra_type = 'subtitle_extra',
                                replacements = replacements,
                                folder_name = folder_name,
                                file_name = file_name,
                                destination = destination,
                                group = group,
                                current_file = current_file
                            )

                            # Don't add language if multiple languages in 1 file
                            if len(sub_langs) > 1:
                                rename_files[current_file] = os.path.join(destination, final_folder_name, final_file_name)
                            elif len(sub_langs) == 1:
                                sub_name = final_file_name.replace(replacements['ext'], '%s.%s' % (sub_langs[0], replacements['ext']))
                                rename_files[current_file] = os.path.join(destination, final_folder_name, sub_name)

                            rename_files = mergeDicts(rename_files, rename_extras)

                        # Filename without cd etc
                        if file_type is 'movie':
                            rename_extras = self.getRenameExtras(
                                extra_type = 'movie_extra',
                                replacements = replacements,
                                folder_name = folder_name,
                                file_name = file_name,
                                destination = destination,
                                group = group,
                                current_file = current_file
                            )
                            rename_files = mergeDicts(rename_files, rename_extras)

                            group['destination_dir'] = os.path.join(destination, final_folder_name)

                        if multiple:
                            cd += 1

                # Before renaming, remove the lower quality files

                library = db.query(Library).filter_by(identifier = group['library']['identifier']).first()
                done_status = fireEvent('status.get', 'done', single = True)
                active_status = fireEvent('status.get', 'active', single = True)

                for movie in library.movies:

                    # Mark movie "done" onces it found the quality with the finish check
                    try:
                        if movie.status_id == active_status.get('id'):
                            for profile_type in movie.profile.types:
                                if profile_type.quality_id == group['meta_data']['quality']['id'] and profile_type.finish:
                                    movie.status_id = done_status.get('id')
                                    db.commit()
                    except Exception, e:
                        log.error('Failed marking movie finished: %s %s' % (e, traceback.format_exc()))

                    # Go over current movie releases
                    for release in movie.releases:

                        # When a release already exists
                        if release.status_id is done_status.get('id'):

                            # This is where CP removes older, lesser quality releases
                            if release.quality.order > group['meta_data']['quality']['order']:
                                log.info('Removing lesser quality %s for %s.' % (movie.library.titles[0].title, release.quality.label))
                                for current_file in release.files:
                                    remove_files.append(current_file)
                                remove_releases.append(release)
                            # Same quality, but still downloaded, so maybe repack/proper/unrated/directors cut etc
                            elif release.quality.order is group['meta_data']['quality']['order']:
                                log.info('Same quality release already exists for %s, with quality %s. Assuming repack.' % (movie.library.titles[0].title, release.quality.label))
                                for current_file in release.files:
                                    remove_files.append(current_file)
                                remove_releases.append(release)

                            # Downloaded a lower quality, rename the newly downloaded files/folder to exclude them from scan
                            else:
                                log.info('Better quality release already exists for %s, with quality %s' % (movie.library.titles[0].title, release.quality.label))

                                # Add _EXISTS_ to the parent dir
                                if group['dirname']:
                                    for rename_me in rename_files: # Don't rename anything in this group
                                        rename_files[rename_me] = None
                                    rename_files[group['parentdir']] = group['parentdir'].replace(group['dirname'], '_EXISTS_%s' % group['dirname'])
                                else: # Add it to filename
                                    for rename_me in rename_files:
                                        filename = os.path.basename(rename_me)
                                        rename_files[rename_me] = rename_me.replace(filename, '_EXISTS_%s' % filename)

                                # Notify on rename fail
                                download_message = 'Renaming of %s (%s) canceled, exists in %s already.' % (movie.library.titles[0].title, group['meta_data']['quality']['label'], release.quality.label)
                                fireEvent('movie.renaming.canceled', message = download_message, data = group)

                                break

                # Remove leftover files
                if self.conf('cleanup') and not self.conf('move_leftover'):
                    log.debug('Removing leftover files')
                    for current_file in group['files']['leftover']:
                        remove_files.append(current_file)

            # Rename all files marked
            group['renamed_files'] = []
            for src in rename_files:
                if rename_files[src]:
                    dst = rename_files[src]
                    log.info('Renaming "%s" to "%s"' % (src, dst))

                    # Create dir
                    self.makeDir(os.path.dirname(dst))

                    try:
                        self.moveFile(src, dst)
                        group['renamed_files'].append(dst)
                    except:
                        log.error('Failed moving the file "%s" : %s' % (os.path.basename(src), traceback.format_exc()))

            # Remove files
            for src in remove_files:

                if isinstance(src, File):
                    src = src.path

                log.info('Removing "%s"' % src)
                try:
                    os.remove(src)
                except:
                    log.error('Failed removing %s: %s' % (src, traceback.format_exc()))

            # Remove matching releases
            for release in remove_releases:
                log.debug('Removing release %s' % release.identifier)
                try:
                    db.delete(release)
                except:
                    log.error('Failed removing %s: %s' % (release.identifier, traceback.format_exc()))

            if group['dirname'] and group['parentdir']:
                try:
                    log.info('Deleting folder: %s' % group['parentdir'])
                    self.deleteEmptyFolder(group['parentdir'])
                except:
                    log.error('Failed removing %s: %s' % (group['parentdir'], traceback.format_exc()))

            # Search for trailers etc
            fireEventAsync('renamer.after', group)

            # Notify on download
            download_message = 'Downloaded %s (%s)' % (group['library']['titles'][0]['title'], replacements['quality'])
            fireEventAsync('movie.downloaded', message = download_message, data = group)

            # Break if CP wants to shut down
            if self.shuttingDown():
                break

        self.renaming_started = False

    def getRenameExtras(self, extra_type = '', replacements = {}, folder_name = '', file_name = '', destination = '', group = {}, current_file = ''):

        rename_files = {}

        def test(s):
            return current_file[:-len(replacements['ext'])] in s

        for extra in set(filter(test, group['files'][extra_type])):
            replacements['ext'] = getExt(extra)

            final_folder_name = self.doReplace(folder_name, replacements)
            final_file_name = self.doReplace(file_name, replacements)
            rename_files[extra] = os.path.join(destination, final_folder_name, final_file_name)

        return rename_files

    def moveFile(self, old, dest):
        try:
            shutil.move(old, dest)

            try:
                os.chmod(dest, Env.getPermission('folder'))
            except:
                log.error('Failed setting permissions for file: %s' % dest)

        except:
            log.error("Couldn't move file '%s' to '%s': %s" % (old, dest, traceback.format_exc()))
            raise Exception

        return True

    def doReplace(self, string, replacements):
        '''
        replace confignames with the real thing
        '''

        replaced = toUnicode(string)
        for x, r in replacements.iteritems():
            if r is not None:
                replaced = replaced.replace('<%s>' % toUnicode(x), toUnicode(r))
            else:
                #If information is not available, we don't want the tag in the filename
                replaced = replaced.replace('<' + x + '>', '')

        replaced = re.sub(r"[\x00:\*\?\"<>\|]", '', replaced)

        sep = self.conf('separator')
        return self.replaceDoubles(replaced).replace(' ', ' ' if not sep else sep)

    def replaceDoubles(self, string):
        return string.replace('  ', ' ').replace(' .', '.')

    def deleteEmptyFolder(self, folder):

        for root, dirs, files in os.walk(folder):

            for dir_name in dirs:
                full_path = os.path.join(root, dir_name)
                if len(os.listdir(full_path)) == 0:
                    try:
                        os.rmdir(full_path)
                    except:
                        log.error('Couldn\'t remove empty directory %s: %s' % (full_path, traceback.format_exc()))

        try:
            os.rmdir(folder)
        except:
            log.error('Couldn\'t remove empty directory %s: %s' % (folder, traceback.format_exc()))