1#!/usr/local/bin/python3.8
2# vim:fileencoding=utf-8
3
4
5__license__ = 'GPL v3'
6__copyright__ = '2015, Kovid Goyal <kovid at kovidgoyal.net>'
7
8import os, sys, time, traceback
9from threading import Thread
10
11
12from calibre import guess_type, prints
13from calibre.constants import is64bit, isportable, isfrozen, __version__, DEBUG
14from calibre.utils.winreg.lib import Key, HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE
15from calibre.utils.lock import singleinstance
16from polyglot.builtins import iteritems, itervalues
17from calibre_extensions import winutil
18
19# See https://msdn.microsoft.com/en-us/library/windows/desktop/cc144154(v=vs.85).aspx
20
21
22def default_programs():
23    return {
24        'calibre.exe': {
25            'icon_id':'main_icon',
26            'description': _('The main calibre program, used to manage your collection of e-books'),
27            'capability_name': 'calibre' + ('64bit' if is64bit else ''),
28            'name': 'calibre' + (' 64-bit' if is64bit else ''),
29            'assoc_name': 'calibre' + ('64bit' if is64bit else ''),
30        },
31
32        'ebook-edit.exe': {
33            'icon_id':'editor_icon',
34            'description': _('The calibre E-book editor. It can be used to edit common e-book formats.'),
35            'capability_name': 'Editor' + ('64bit' if is64bit else ''),
36            'name': 'calibre Editor' + (' 64-bit' if is64bit else ''),
37            'assoc_name': 'calibreEditor' + ('64bit' if is64bit else ''),
38        },
39
40        'ebook-viewer.exe': {
41            'icon_id':'viewer_icon',
42            'description': _('The calibre E-book viewer. It can view most known e-book formats.'),
43            'capability_name': 'Viewer' + ('64bit' if is64bit else ''),
44            'name': 'calibre Viewer' + (' 64-bit' if is64bit else ''),
45            'assoc_name': 'calibreViewer' + ('64bit' if is64bit else ''),
46        },
47    }
48
49
50def extensions(basename):
51    if basename == 'calibre.exe':
52        from calibre.ebooks import BOOK_EXTENSIONS
53        # We remove rar and zip as they interfere with 7-zip associations
54        # https://www.mobileread.com/forums/showthread.php?t=256459
55        return set(BOOK_EXTENSIONS) - {'rar', 'zip'}
56    if basename == 'ebook-viewer.exe':
57        from calibre.customize.ui import all_input_formats
58        return set(all_input_formats())
59    if basename == 'ebook-edit.exe':
60        from calibre.ebooks.oeb.polish.main import SUPPORTED
61        from calibre.ebooks.oeb.polish.import_book import IMPORTABLE
62        return SUPPORTED | IMPORTABLE
63
64
65class NotAllowed(ValueError):
66    pass
67
68
69def check_allowed():
70    if not isfrozen:
71        raise NotAllowed('Not allowed to create associations for non-frozen installs')
72    if isportable:
73        raise NotAllowed('Not allowed to create associations for portable installs')
74    if sys.getwindowsversion()[:2] < (6, 2):
75        raise NotAllowed('Not allowed to create associations for windows versions older than Windows 8')
76    if 'CALIBRE_NO_DEFAULT_PROGRAMS' in os.environ:
77        raise NotAllowed('Disabled by the CALIBRE_NO_DEFAULT_PROGRAMS environment variable')
78
79
80def create_prog_id(ext, prog_id, ext_map, exe):
81    with Key(r'Software\Classes\%s' % prog_id) as key:
82        type_name = _('%s Document') % ext.upper()
83        key.set(value=type_name)
84        key.set('FriendlyTypeName', type_name)
85        key.set('PerceivedType', 'Document')
86        key.set(sub_key='DefaultIcon', value=exe+',0')
87        key.set_default_value(r'shell\open\command', '"%s" "%%1"' % exe)
88        # contrary to the msdn docs, this key prevents calibre programs
89        # from appearing in the initial open with list, see
90        # https://www.mobileread.com/forums/showthread.php?t=313668
91        # key.set('AllowSilentDefaultTakeOver')
92
93    with Key(r'Software\Classes\.%s\OpenWithProgIDs' % ext) as key:
94        key.set(prog_id)
95
96
97def progid_name(assoc_name, ext):
98    return '%s.AssocFile.%s' % (assoc_name, ext.upper())
99
100
101def cap_path(data):
102    return r'Software\calibre\%s\Capabilities' % data['capability_name']
103
104
105def register():
106    base = os.path.dirname(sys.executable)
107
108    for program, data in iteritems(default_programs()):
109        data = data.copy()
110        exe = os.path.join(base, program)
111        capabilities_path = cap_path(data)
112        ext_map = {ext.lower():guess_type('file.' + ext.lower())[0] for ext in extensions(program)}
113        ext_map = {ext:mt for ext, mt in iteritems(ext_map) if mt}
114        prog_id_map = {ext:progid_name(data['assoc_name'], ext) for ext in ext_map}
115
116        with Key(capabilities_path) as key:
117            for k, v in iteritems({'ApplicationDescription':'description', 'ApplicationName':'name'}):
118                key.set(k, data[v])
119            key.set('ApplicationIcon', '%s,0' % exe)
120            key.set_default_value(r'shell\open\command', '"%s" "%%1"' % exe)
121
122            with Key('FileAssociations', root=key) as fak, Key('MimeAssociations', root=key) as mak:
123                # previous_associations = set(fak.values())
124                for ext, prog_id in iteritems(prog_id_map):
125                    mt = ext_map[ext]
126                    fak.set('.' + ext, prog_id)
127                    mak.set(mt, prog_id)
128        for ext, prog_id in iteritems(prog_id_map):
129            create_prog_id(ext, prog_id, ext_map, exe)
130
131        with Key(r'Software\RegisteredApplications') as key:
132            key.set(data['name'], capabilities_path)
133
134    winutil.notify_associations_changed()
135
136
137def unregister():
138    for program, data in iteritems(default_programs()):
139        capabilities_path = cap_path(data).rpartition('\\')[0]
140        ext_map = {ext.lower():guess_type('file.' + ext.lower())[0] for ext in extensions(program)}
141        ext_map = {ext:mt for ext, mt in iteritems(ext_map) if mt}
142        prog_id_map = {ext:progid_name(data['assoc_name'], ext) for ext in ext_map}
143        with Key(r'Software\RegisteredApplications') as key:
144            key.delete_value(data['name'])
145        parent, sk = capabilities_path.rpartition('\\')[0::2]
146        with Key(parent) as key:
147            key.delete_tree(sk)
148        for ext, prog_id in iteritems(prog_id_map):
149            with Key(r'Software\Classes\.%s\OpenWithProgIDs' % ext) as key:
150                key.delete_value(prog_id)
151            with Key(r'Software\Classes') as key:
152                key.delete_tree(prog_id)
153
154
155class Register(Thread):
156
157    daemon = True
158
159    def __init__(self, prefs):
160        Thread.__init__(self, name='RegisterDP')
161        self.prefs = prefs
162        self.start()
163
164    def run(self):
165        try:
166            self.do_register()
167        except Exception:
168            traceback.print_exc()
169
170    def do_register(self):
171        try:
172            check_allowed()
173        except NotAllowed:
174            return
175        if singleinstance('register_default_programs'):
176            if self.prefs.get('windows_register_default_programs', None) != __version__:
177                self.prefs['windows_register_default_programs'] = __version__
178                if DEBUG:
179                    st = time.monotonic()
180                    prints('Registering with default programs...')
181                register()
182                if DEBUG:
183                    prints('Registered with default programs in %.1f seconds' % (time.monotonic() - st))
184
185    def __enter__(self):
186        return self
187
188    def __exit__(self, *args):
189        # Give the thread some time to finish in case the user quit the
190        # application very quickly
191        self.join(4.0)
192
193
194def get_prog_id_map(base, key_path):
195    desc, ans = None, {}
196    try:
197        k = Key(open_at=key_path, root=base)
198    except OSError as err:
199        if err.winerror == winutil.ERROR_FILE_NOT_FOUND:
200            return desc, ans
201        raise
202    with k:
203        desc = k.get_mui_string('ApplicationDescription')
204        if desc is None:
205            return desc, ans
206        for ext, prog_id in k.values(sub_key='FileAssociations', get_data=True):
207            ans[ext[1:].lower()] = prog_id
208    return desc, ans
209
210
211def get_open_data(base, prog_id):
212    try:
213        k = Key(open_at=r'Software\Classes\%s' % prog_id, root=base)
214    except OSError as err:
215        if err.winerror == winutil.ERROR_FILE_NOT_FOUND:
216            return None, None, None
217    with k:
218        cmd = k.get(sub_key=r'shell\open\command')
219        if cmd:
220            parts = cmd.split()
221            if parts[-1] == '/dde' and '%1' not in cmd:
222                cmd = ' '.join(parts[:-1]) + ' "%1"'
223        return cmd, k.get(sub_key='DefaultIcon'), k.get_mui_string('FriendlyTypeName') or k.get()
224
225
226def split_commandline(commandline):
227    # CommandLineToArgvW returns path to executable if called with empty string.
228    if not commandline.strip():
229        return []
230    return list(winutil.parse_cmdline(commandline))
231
232
233def friendly_app_name(prog_id=None, exe=None):
234    try:
235        return winutil.friendly_name(prog_id, exe)
236    except Exception:
237        traceback.print_exc()
238
239
240def find_programs(extensions):
241    extensions = frozenset(extensions)
242    ans = []
243    seen_prog_ids, seen_cmdlines = set(), set()
244
245    # Search for programs registered using Default Programs that claim they are
246    # capable of handling the specified extensions.
247    for base in (HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE):
248        try:
249            k = Key(open_at=r'Software\RegisteredApplications', root=base)
250        except OSError as err:
251            if err.winerror == winutil.ERROR_FILE_NOT_FOUND:
252                continue
253            raise
254        with k:
255            for name, key_path in k.values(get_data=True):
256                try:
257                    app_desc, prog_id_map = get_prog_id_map(base, key_path)
258                except Exception:
259                    traceback.print_exc()
260                    continue
261                for ext in extensions:
262                    prog_id = prog_id_map.get(ext)
263                    if prog_id is not None and prog_id not in seen_prog_ids:
264                        seen_prog_ids.add(prog_id)
265                        cmdline, icon_resource, friendly_name = get_open_data(base, prog_id)
266                        if cmdline and cmdline not in seen_cmdlines:
267                            seen_cmdlines.add(cmdline)
268                            ans.append({'name':app_desc, 'cmdline':cmdline, 'icon_resource':icon_resource})
269
270    # Now look for programs that only register with Windows Explorer instead of
271    # Default Programs (for example, FoxIt PDF reader)
272    for ext in extensions:
273        try:
274            k = Key(open_at=r'Software\Microsoft\Windows\CurrentVersion\Explorer\FileExts\.%s\OpenWithProgIDs' % ext, root=HKEY_CURRENT_USER)
275        except OSError as err:
276            if err.winerror == winutil.ERROR_FILE_NOT_FOUND:
277                continue
278        for prog_id in itervalues(k):
279            if prog_id and prog_id not in seen_prog_ids:
280                seen_prog_ids.add(prog_id)
281                cmdline, icon_resource, friendly_name = get_open_data(base, prog_id)
282                if cmdline and cmdline not in seen_cmdlines:
283                    seen_cmdlines.add(cmdline)
284                    exe_name = None
285                    exe = split_commandline(cmdline)
286                    if exe:
287                        exe_name = friendly_app_name(prog_id) or os.path.splitext(os.path.basename(exe[0]))[0]
288                    name = exe_name or friendly_name
289                    if name:
290                        ans.append({'name':name, 'cmdline':cmdline, 'icon_resource':icon_resource})
291    return ans
292
293
294if __name__ == '__main__':
295    from pprint import pprint
296    pprint(find_programs('docx'.split()))
297