make download finder more intelligent

1 files changed, 47 insertions, 21 deletions

diff --git a/update.py b/update.py
index a7d3ab2..38af029 100755
--- a/update.py
+++ b/update.py
@@ -6,6 +6,8 @@ import sys
 import hashlib
 import shutil
 import re
+import collections
+import urllib.parse
 
 import requests
 
@@ -38,6 +40,10 @@ parser.add_argument('--whitelist-file',
     type=str,
     default="whitelist.txt",
     help="Optional custom whitelist file that tells 'install' which files not to remove (default: whitelist.txt)")
+parser.add_argument("--game-version",
+    type=str,
+    default=None,
+    help="The maximum game version to update mods to")
 
 ## loaded from version.txt
 VERSION = 0
@@ -48,9 +54,9 @@ def read_file(fil):
         for line in f:
             string = line.strip().split()
             if len(line) > 1 and line[0] != '#':
-                # run strip on each element
-                string = tuple(map(lambda x: x.strip(), string))
-                strings.append(string)
+                    # run strip on each element
+                    string = tuple(map(lambda x: x.strip(), string))
+                    strings.append(string)
 
     return strings
 
@@ -92,6 +98,10 @@ def install(args):
 
 # Using the latest urls, update downloads.txt to match and have the correct sha1
 def apply_updates(args):
+    if args.game_version is not None:
+        version = tuple(int(x) for x in args.game_version.split('.'))
+    else:
+        version = (2, 0, 0)
     print("Populating version File...")
     mods = read_file(args.filename)
     print("Getting new versions of all mods...")
@@ -102,7 +112,7 @@ def apply_updates(args):
         for mod in mods:
             print("Fetching {mod[0]}...".format(mod=mod))
             if 'curseforge' in mod[1]:
-                url = find_cdn(ffx, mod[1])
+                url = find_cdn(ffx, mod[1], version)
             else:
                 url = requests.get(mod[1]).url
             if url is None:
@@ -120,6 +130,10 @@ def apply_updates(args):
 
 # Find if any updates are available
 def check_updates(args):
+    if args.game_version is not None:
+        version = tuple(int(x) for x in args.game_version.split('.'))
+    else:
+        version = (2, 0, 0)
     print("Checking for updates to version " + str(VERSION) + "...")
     latest = read_file(args.filename)
     old = read_file(args.version_file)
@@ -133,7 +147,7 @@ def check_updates(args):
         print("Checking for updates to {mod[0]}...".format(mod=mod), end="")
         sys.stdout.flush()  # takes care of line-buffered terminals
         if 'curseforge' in mod[1]:
-            url = find_cdn(ffx, mod[1])
+            url = find_cdn(ffx, mod[1], version)
         else:
             url = requests.get(mod[1]).url
         if url in old_urls:
@@ -147,27 +161,39 @@ def check_updates(args):
     if num_updates >= 0:
         print("Run 'python update.py apply_updates' to create a new version with these updates applied.")
 
-# Use selenium to find curseforge CDN links around cloudflare
-def find_cdn(ffx, url):
+
+def find_cdn(ffx, url, version):
+    """
+    Given a mod home URL, finds the most up-to-date mod version compatible with the given game version.
+    Returns the direct Forge CDN download URL
+    """
+    #TODO filter mods by forge/fabric compatibility
     try:
-        #ffx.get(url + '/download')
-        ffx.get(url + '/files')
-        #page_src = ffx.page_source
-        #dl = re.search('Elerium.PublicProjectDownload.countdown\(".*?"\);', page_src)
-        #if not dl:
-        #    return None
-        dl = ffx.find_element_by_xpath("html/body/div/main/div/div/section/div/div/div/section/article/div/div/a").get_attribute("href")
-        dl = re.search('\d{7}', dl)
-        dl = dl.group(0)
-        four = str(int(dl[:4]))
-        three = str(int(dl[4:]))
-
-        file_name = ffx.find_elements_by_xpath("html/body/div/main/div/div/section/div/div/div/section/article/div/div/span[contains(@class, 'text-sm')]")[1].text
-        return 'https://media.forgecdn.net/files/{four}/{three}/{jar}'.format(four=four,three=three,jar=file_name)
+        ffx.get(url + '/files/all')
+        mod_versions = ffx.find_elements_by_class_name("listing")[0].find_elements_by_xpath("tbody/tr") # extract the table of files from the page
+        row_info = collections.namedtuple("row_info", ["type", "filename", "cdn_id", "game_version"]) # create a custom tuple because data
+        rows = []
+        for version_entry in mod_versions:
+            # parse out the four fields that we use
+            entry_cells = version_entry.find_elements_by_tag_name("td")
+            release_type = entry_cells[0].text
+            filename = urllib.parse.quote(entry_cells[1].find_elements_by_tag_name("a")[0].text)
+            try:
+                game_version = tuple([int(x) for x in entry_cells[4].find_element_by_class_name("mr-2").text.split(".")]) # get game version and convert to tuple
+            except:
+                game_version = (0, 0, 0)
+            cdn_id = entry_cells[1].find_element_by_tag_name("a").get_property("href").split("/")[-1]
+            rows.append(row_info(release_type, filename, cdn_id, game_version))
+        rows.sort(key=lambda x: x.game_version, reverse=True)
+        best_row = next(x for x in rows if x.game_version <= version)
+
+        return f'https://media.forgecdn.net/files/{best_row.cdn_id[:4]}/{best_row.cdn_id[4:]}/{best_row.filename}'
 
     except:
+        import traceback; traceback.print_exc()
         return None
 
+
 def firefox():
     print("Starting Selenium...")
     try: