1import os
2import xml.etree.ElementTree as ET
3
4from loguru import logger
5
6from flexget import plugin
7from flexget.event import event
8
9try:
10    # NOTE: Importing other plugins is discouraged!
11    from flexget.components.imdb.utils import is_valid_imdb_title_id
12except ImportError:
13    raise plugin.DependencyError(issued_by=__name__, missing='imdb')
14
15
16logger = logger.bind(name='nfo_lookup')
17
18
19class NfoLookup:
20    """
21    Retrieves information from a local '.nfo' info file.
22
23    The read metadata will be add as 'nfo_something' in the entry. Also, if an 'id' is found in the '.nfo' file then the
24    'imdb_id' field will be set to its value. This means that if the imdb_lookup plugin is used in addition to this
25    plugin it will be able to use the ID from '.nfo' file to get the correct movie.
26
27    The nfo file is used by Kodi.
28
29    Example:
30        nfo_lookup: yes
31
32    WARNING: This plugin will read a file with extension '.nfo' and the same name as the entry filename as an XML file
33    using xml.etree.ElementTree from the standard python library. As such, it is vulnerable to XML vulnerabilities
34    described in the link below
35    https://docs.python.org/3/library/xml.html#xml-vulnerabilities
36
37    Use this only with nfo files you have created yourself.
38    """
39
40    schema = {'type': 'boolean'}
41    nfo_file_extension = '.nfo'
42
43    # This priority makes sure this plugin runs before the imdb_lookup plugin, if it is also used. That way setting
44    # imdb_id here will help imdb_lookup find the correct movie.
45    @plugin.priority(150)
46    def on_task_metainfo(self, task, config):
47        # check if disabled (value set to false)
48        if not config:
49            # Config was set to 'no' instead of yes. Don't do anything then.
50            return
51
52        for entry in task.entries:
53            # If this entry was obtained from the filesystem plugin it should have a filename field. If it does not have
54            # one then there is nothing we can do in this plugin.
55            filename = entry.get('filename')
56            location = entry.get('location')
57
58            # If there is no 'filename' field there is also no nfo file
59            if filename is None or location is None:
60                logger.warning(
61                    "Entry {} didn't come from the filesystem plugin", entry.get('title')
62                )
63                continue
64            else:
65                # This will be None if there is no nfo file
66                nfo_filename = self.get_nfo_filename(entry)
67                if nfo_filename is None:
68                    logger.warning(
69                        'Entry {} has no corresponding {} file',
70                        entry.get('title'),
71                        self.nfo_file_extension,
72                    )
73                    continue
74
75            # Populate the fields from the information in the .nfo file Note that at this point `nfo_filename` has the
76            # name of an existing '.nfo' file
77            self.lookup(entry, nfo_filename)
78
79    def lookup(self, entry, nfo_filename):
80        # If there is already data from a previous parse then we don't need to do anything
81        if entry.get('nfo_id') is not None:
82            logger.warning(
83                'Entry {} was already parsed by nfo_lookup and it will be skipped. ',
84                entry.get('title'),
85            )
86            return
87
88        # nfo_filename Should not be None at this point
89        assert nfo_filename is not None
90
91        # Get all values we can from the nfo file. If the nfo file can't be parsed then a warning is logged and we
92        # return without changing the entry
93        try:
94            nfo_reader = NfoReader(nfo_filename)
95            fields = nfo_reader.get_fields_from_nfo_file()
96        except BadXmlFile:
97            logger.warning("Invalid '.nfo' file for entry {}", entry.get('title'))
98            return
99
100        entry.update(fields)
101
102        # If a valid IMDB id was found in the nfo file, set the imdb_id field of the entry. This will help the
103        # imdb_lookup plugin to get the correct data if it is also used.
104        if 'nfo_id' in fields:
105            if is_valid_imdb_title_id(entry.get('nfo_id', '')):
106                entry.update({'imdb_id': fields['nfo_id']})
107            else:
108                logger.warning(
109                    "ID found in nfo file for entry '{}', but it was not a valid IMDB ID",
110                    entry.get('title'),
111                )
112
113    def get_nfo_filename(self, entry):
114        """
115        Get the filename of the nfo file from the 'location' in the entry.
116
117        Returns
118        -------
119        str
120            The file name of the 'nfo' file, or None it there is no 'nfo' file.
121        """
122        location = entry.get('location')
123        nfo_full_filename = os.path.splitext(location)[0] + self.nfo_file_extension
124
125        if os.path.isfile(nfo_full_filename):
126            return nfo_full_filename
127
128
129class BadXmlFile(Exception):
130    """
131    Exception that is raised if the nfo file can't be parsed due to some invalid nfo file.
132    """
133
134    pass
135
136
137class NfoReader:
138    """
139    Class in charge of parsing the '.nfo' file and getting a dictionary of fields.
140
141    The '.nfo' file is an XML file. Some fields can only appear once, such as 'title', 'id', 'plot', etc., while other
142    fields can appear multiple times (with different values), such as 'thumb', 'genre', etc. These fields are listed in
143    the `_fields` attribute.
144    """
145
146    def __init__(self, filename):
147        try:
148            tree = ET.parse(filename)
149            root = tree.getroot()
150        except ET.ParseError:
151            raise BadXmlFile()
152
153        if os.path.exists(filename):
154            self._nfo_filename = filename
155            self._root = root
156        else:
157            raise BadXmlFile()
158
159        # Each key in the dictionary correspond to a field that should be read from the nfo file. The values are a tuple
160        # with a boolean and a callable. The boolean indicates if the field can appear multiple times, while the
161        # callable is a function to read the field value from the XML element.
162        #
163        # In the future we could extend the nfo_lookup plugin to accept 'set' in its configuration to add new entries to
164        # this dictionary to handle other tags in the nfo file and add the data to the entry.
165        self._fields = {
166            "title": (False, NfoReader._single_elem_getter_func),
167            "originaltitle": (False, NfoReader._single_elem_getter_func),
168            "sorttitle": (False, NfoReader._single_elem_getter_func),
169            "rating": (False, NfoReader._single_elem_getter_func),
170            "year": (False, NfoReader._single_elem_getter_func),
171            "votes": (False, NfoReader._single_elem_getter_func),
172            "plot": (False, NfoReader._single_elem_getter_func),
173            "runtime": (False, NfoReader._single_elem_getter_func),
174            "id": (False, NfoReader._single_elem_getter_func),
175            "filenameandpath": (False, NfoReader._single_elem_getter_func),
176            "trailer": (False, NfoReader._single_elem_getter_func),
177            "thumb": (True, NfoReader._single_elem_getter_func),
178            "genre": (True, NfoReader._single_elem_getter_func),
179            "director": (True, NfoReader._single_elem_getter_func),
180            # Actor field has child elements, such as 'name' and 'role'
181            "actor": (True, NfoReader._composite_elem_getter_func),
182            "studio": (True, NfoReader._single_elem_getter_func),
183            "country": (True, NfoReader._single_elem_getter_func),
184        }
185
186    @staticmethod
187    def _single_elem_getter_func(x):
188        """
189        Method to get the text value of simple XML element that does not contain child nodes.
190        """
191        return x.text
192
193    @staticmethod
194    def _composite_elem_getter_func(x):
195        """
196        Method to get XML elements that have children as a dictionary.
197        """
198        return {i.tag: i.text for i in x}
199
200    def _extract_single_field(self, name, getter_func):
201        """
202        Use this method to get fields from the root XML tree that only appear once, such as 'title', 'year', etc.
203        """
204        f = self._root.find(name)
205        if f is not None:
206            return getter_func(f)
207
208    def _extract_multiple_field(self, name, getter_func):
209        """
210        Use this method to get fields from the root XML tree that can appear more than once, such as 'actor', 'genre',
211        'director', etc. The result will be a list of values.
212        """
213        values = [getter_func(i) for i in self._root.findall(name)]
214
215        if len(values) > 0:
216            return values
217
218    def get_fields_from_nfo_file(self):
219        """
220        Returns a dictionary with all firlds read from the '.nfo' file.
221
222        The keys are named as 'nfo_something'.
223        """
224        d = {}
225        if self._root is None:
226            return d
227
228        # TODO: Right now it only works for movies
229        if self._root.tag != 'movie':
230            return d
231
232        for name, values in self._fields.items():
233            multiple_bool = values[0]
234            getter_func = values[1]
235
236            nfo_field_name = 'nfo_{0}'.format(name)
237
238            if multiple_bool:
239                v = self._extract_multiple_field(name, getter_func)
240            else:
241                v = self._extract_single_field(name, getter_func)
242
243            if v is not None:
244                d[nfo_field_name] = v
245
246        return d
247
248
249@event('plugin.register')
250def register_plugin():
251    plugin.register(NfoLookup, 'nfo_lookup', api_ver=2)
252