Fix #10 as well as some weird behaviour

decorators.py

@@ -7,9 +7,13 @@
 def box_on_error(title):
     def decorator(func):
         @functools.wraps(func)
-        def wrapper():
+        def wrapper(*args, **kwargs):
+            print("args: " + str(args))
+
             try:
-                return func()
+                if isinstance(args[0], bool):
+                    return func(*args[1:], **kwargs)
+                return func(*args, **kwargs)
             except Exception as e:
                 box = QtWidgets.QMessageBox()
                 box.setWindowTitle(title)

serverIndex.py

@@ -34,8 +34,11 @@ def setup_index_list(ui, options: Options):
     for entry in index:
         ui.indexList.addTopLevelItem(QtWidgets.QTreeWidgetItem([entry['url'], "Yes" if entry['wls'] else "No", "Pinging..."]))
         # Start a thread to ping each server and update the list
-        def ping_server(item):
+        def ping_server(item: QtWidgets.QTreeWidgetItem):
             pingval = ping(item.text(0))
+            # Make sure item still exists
+            if not item:
+                return
             if pingval == -1:
                 item.setText(2, "Offline")
             else:

vrmlUpdater.py

@@ -1,28 +1,42 @@
 import os
 import re
+from decorators import box_on_error
 
+@box_on_error("Error updating wrl list")
 def set_server(path: str, ip: str, wls: bool) -> None:
     print("Setting server for " + path + " to " + ip)
+    updated = False
     with open(path, 'r') as f:
         # Create a temporary file to store the new .wrl
         with open('temp.wrl', 'w') as temp:
             # Locate and empty every field cpBureau and cpBureauWLS
-            try:
-                for line in f:
-                    if 'cpBureau' or 'cpBureauWLS' in line:
-                        line = re.sub(r'(cpBureau(WLS)?\s*)".*"', r'\1""', line)
+            for line in f:
+                if 'cpBureau' or 'cpBureauWLS' in line:
+                    line = re.sub(r'(cpBureau(WLS)?\s*)".*"', r'\1""', line)
+
+                # If we are using WLS, replace the WLS field with the ip
+                if wls and 'cpBureauWLS' in line:
+                    line = re.sub(r'(cpBureauWLS\s*)".*"', r'\1"{}"'.format(ip), line)
+                    updated = True    
+                # If we are not using WLS, replace the non-WLS field with the ip
+                elif not wls and 'cpBureau' in line:
+                    line = re.sub(r'(cpBureau\s*)".*"', r'\1"{}"'.format(ip), line)
+                    updated = True
+
+                temp.write(line)
+
+            if not updated:
+                # We need to add the field in the correct place in the file
+                temp.write('Sony_WorldInfo {\n')
+                if wls:
+                    temp.write('\tcpBureauWLS "{}"\n'.format(ip))
+                    temp.write('\tcpBureau ""\n')
+                else:
+                    temp.write('\tcpBureauWLS ""\n')
+                    temp.write('\tcpBureau "{}"\n'.format(ip))
+                temp.write('}\n')
 
-                    # If we are using WLS, replace the WLS field with the ip
-                    if wls and 'cpBureauWLS' in line:
-                            line = re.sub(r'(cpBureauWLS\s*)".*"', r'\1"{}"'.format(ip), line)
-                    # If we are not using WLS, replace the non-WLS field with the ip
-                    elif not wls and 'cpBureau' in line:
-                        line = re.sub(r'(cpBureau\s*)".*"', r'\1"{}"'.format(ip), line)
-
-                    temp.write(line)
-            except:
-                print("Error setting server for " + path)
-                return
+
 
     # Replace the original .wrl with the new .wrl
     os.rename(path, path + ".OLD")

wrls.py

@@ -21,7 +21,7 @@ def __init__(self, options: Options):
                     break
                 if file.endswith(".wrl"):
                     try:
-                        if "cpBureau" in open(os.path.join(root, file)).read():
+                        if "WorldInfo" in open(os.path.join(root, file)).read():
                          self.wrls.append(os.path.join(root, file))
                     except:
                         pass