[Next] grouped-window-list: Add PID validation in utils.py (hardening improvement)

files/usr/share/cinnamon/applets/grouped-window-list@cinnamon.org/utils.py

@@ -22,91 +22,91 @@ def spawn(command):
     out = process.stdout.decode('utf-8')
     return out
 
+def get_valid_pid():
+    if len(CLI) < 3 or not CLI[2].isdigit():
+        print("Invalid PID provided")
+        sys.exit(1)
+    return int(CLI[2])
+
+def read_process_cmdline(pid):
+    return spawn(f"cat /proc/{pid}/cmdline")
+
+def normalize_process(process):
+    if '.exe' in process:
+        if 'Z:' in process:
+            process = process.split('Z:')[1]
+        process = process.replace('\\', '/')
+        process = process.split('.exe')[0] + '.exe'
+        process = 'wine ' + process.replace(' ', r'\ ')
+    else:
+        process = process[:-1]  # remove trailing null byte
+
+    if process == 'python mainwindow.py':
+        process = 'playonlinux'
+
+    return process
+
 # Utility script that creates GDesktop files for Wine and other window backed applications.
 def handle_cli():
+    if len(CLI) < 2 or CLI[1] != 'get_process':
+        return
 
-    if CLI[1] == 'get_process':
-        process = spawn('cat /proc/{}/cmdline'.format(CLI[2]))
+    pid = get_valid_pid()
+    process = read_process_cmdline(pid)
+    process = normalize_process(process)
 
-        if '.exe' in process:
-            if 'Z:' in process:
-                process = process.split('Z:')[1]
+    try:
+        proc_array = process.split('/')
+        process_name = proc_array[-1].title()
 
-            process = process.replace('\\', '/')
-            process = process.split('.exe')[0] + '.exe'
-            process = 'wine '+process.replace(' ', r'\ ')
+        icons_dir = os.path.join(os.getenv('HOME'), '.local/share/icons/hicolor/48x48/apps/')
 
-        process = json.dumps(process)
-        if '\\u0000' in process:
-            process = process.replace('\\u0000', ' ')
-        process = json.loads(process)
+        if '\\ ' in process_name:
+            process_name = process_name.replace('\\ ', ' ')
+        if '.Exe' in process_name:
+            process_name = process_name.replace('.Exe', '')
 
-        if '.exe' not in process:
-            process = process[:-1]
+        icon_file = process_name + '.png'
+        if ' ' in icon_file:
+            icon_file = icon_file.replace(' ', '')
 
-        if process == 'python mainwindow.py':
-            process = 'playonlinux'
+        icon = os.path.join(icons_dir, icon_file)
 
         try:
-            proc_array = process.split('/')
-            pa_len = len(proc_array)
-            process_name = proc_array[pa_len - 1].title()
-
-            # Since this is a window backed app, make sure it has an icon association.
-
-            icons_dir = '{}/.local/share/icons/hicolor/48x48/apps/'.format(os.getenv('HOME'))
-
-            if '\\ ' in process_name:
-                process_name = process_name.replace('\\ ', ' ')
-
-            if '.Exe' in process_name:
-                process_name = process_name.replace('.Exe', '')
+            try:
+                spawn('gnome-exe-thumbnailer {} {}'.format(process.split('wine ')[1], icon))
+            except IndexError:
+                spawn('gnome-exe-thumbnailer {} {}'.format(process, icon))
+        except subprocess.CalledProcessError:
+            icon = None
+
+        g_menu = '[Desktop Entry]\n' \
+                 'Type=Application\n' \
+                 'Encoding=UTF-8\n' \
+                 'Name={}\n' \
+                 'Comment={}\n' \
+                 'Exec={}\n' \
+                 'Terminal=false\n' \
+                 'StartupNotify=true\n'.format(process_name, process_name, process)
+
+        if icon:
+            g_menu += 'Icon={}\n'.format(icon)
 
-            icon_file = process_name+'.png'
+        if '.exe' in process:
+            g_menu += 'GenericName=Wine application\n' \
+                      'Categories=Wine;\n' \
+                      'MimeType=application/x-ms-dos-executable;application/x-msi;application/x-ms-shortcut;\n'
 
-            if ' ' in icon_file:
-                icon_file = icon_file.replace(' ', '')
+        desktop_file = f"{process_name}.cinnamon-generated.desktop"
+        desktop_path = os.path.join(os.getenv('HOME'), '.local/share/applications', desktop_file)
 
-            icon = icons_dir+icon_file
+        with open(desktop_path, 'w', encoding='utf-8') as desktop:
+            print(g_menu)
+            desktop.write(g_menu)
+            spawn(f"chmod +x {desktop_path}")
+            print(desktop_file)
 
-            try:
-                try:
-                    spawn('gnome-exe-thumbnailer {} {}'.format(process.split('wine ')[1], icon))
-                except IndexError:
-                    spawn('gnome-exe-thumbnailer {} {}'.format(process, icon))
-            except subprocess.CalledProcessError:
-                icon = None
-
-            g_menu = '[Desktop Entry]\n' \
-                     'Type=Application\n' \
-                     'Encoding=UTF-8\n' \
-                     'Name={}\n' \
-                     'Comment={}\n' \
-                     'Exec={}\n' \
-                     'Terminal=false\n' \
-                     'StartupNotify=true\n'.format(process_name, process_name, process)
-
-            if icon:
-                g_menu += 'Icon={}\n'.format(icon)
-
-
-            if '.exe' in process:
-                g_menu += 'GenericName=Wine application\n' \
-                          'Categories=Wine;\n' \
-                          'MimeType=application/x-ms-dos-executable;' \
-                          'application/x-msi;application/x-ms-shortcut; \n' \
-
-            desktop_file = '{}.cinnamon-generated.desktop'.format(process_name)
-            desktop_path = '{}/.local/share/applications/{}'.format(os.getenv('HOME'), desktop_file)
-
-            with open(desktop_path, 'w', encoding='utf-8') as desktop:
-                print(g_menu)
-                desktop.write(g_menu)
-                spawn('chmod +x {}'.format(desktop_path))
-                print(desktop_file)
-
-        except KeyError as err:
-            print(err)
-            return
+    except (KeyError, IndexError) as err:
+        print(err)
 
 handle_cli()