import dorkbox.cabParser.CabException;
import dorkbox.cabParser.extractor.CabExtractor;
import org.jsoup.Jsoup;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.nio.file.StandardCopyOption;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.stream.Collectors;

public class DriverDownloader extends Thread {
    private static final String url = "https://www.catalog.update.microsoft.com/DownloadDialog.aspx";
    private static final String formData = "[{\"size\":0,\"languages\":\"\",\"uidInfo\":\"%s\",\"updateID\":\"%s\"}]";
    private static final String downloadString = "downloadInformation[0].files[0].url = '";
    private List<String> driverFilePaths = new ArrayList<>();

    private String[] imports;
    private String uuid;
    private String downloadLink;

    public DriverDownloader(String uuid, String[] imports) {
        this.uuid = uuid;
        this.imports = imports;
    }

    @Override
    public void run() {
        try {
            // query download page
            var result = Jsoup
                    .connect(url)
                    .data("updateIDs", String.format(formData, uuid, uuid))
                    .data("updateIDsBlockedForImport", "")
                    .data("wsusApiPresent", "")
                    .data("contentImport", "")
                    .data("sku", "")
                    .data("serverName", "")
                    .data("ssl", "")
                    .data("portNumber", "")
                    .data("version", "")
                    .post();

            /**
             *  for some reason download link is not in the a tag so i have to parse it out of javascript.
             */
            final var urlIndex = result.toString().indexOf(downloadString) + downloadString.length();
            var possibleUrl = result
                    .toString()
                    .substring(
                            urlIndex,
                            result.toString().indexOf("'", urlIndex));

            if(possibleUrl.endsWith(".cab")) {
                this.downloadLink = possibleUrl;
                downloadCabFile();
                findDriverFiles();
                scanFilesAndDelete();
            }
            else
                this.downloadLink = "none";
        } catch (IOException e) { }
    }

    private void scanFilesAndDelete() {
        driverFilePaths.parallelStream().forEach(path -> {
            System.out.printf("[+] checking driver at path: %s\n", path);
            try {
                /**
                 * run agent-jones.exe to get imports. this is shit code and should be replaced with
                 * a library that can get imports but it seems none are good enough :(
                 */
                Process process = Runtime.getRuntime().exec(new File("agent-jones.exe")
                        .getAbsolutePath()
                        .concat(" ")
                        .concat(new File(path).getAbsolutePath()));
                int ch;
                StringBuilder sb = new StringBuilder();
                while ((ch = process.getInputStream().read()) != -1)
                    sb.append((char) ch);

                var driverImports = sb.toString().split(",");
                if (Arrays.asList(driverImports).containsAll(Arrays.asList(this.imports))) {
                    System.out.println("[+] found with all desired imports....");
                    try {
                        var driverPath = path.split("\\\\");
                        Files.move(
                                new File(path).toPath(),
                                new File("results/"
                                        .concat(uuid)
                                        .concat(driverPath[driverPath.length - 1])) // just in case the driver is in a sub folder
                                        .toPath());
                    } catch (IOException e) { }
                }
            } catch (Exception e) { e.printStackTrace(); }
        });

        /**
         * delete all the files in the unzipped folder + the cab file
         */
        try {
            Files.walk(new File("drivers/".concat(uuid)).toPath()).forEach(file -> new File(file.toString()).delete());
            new File("drivers/".concat(uuid)).delete();
            new File("drivers/".concat(uuid).concat(".cab")).delete();
        } catch (IOException e) { }
    }

    /**
     * scan the unzipped folder for driver files and append them to a list
     */
    private void findDriverFiles() {
        try {
            this.driverFilePaths = Files.walk(Paths.get("drivers/".concat(uuid)))
                    .filter(Files::isRegularFile)
                    .map(x -> x.toString())
                    .collect(Collectors.toList());
            this.driverFilePaths = driverFilePaths.parallelStream().filter(file -> file.endsWith(".sys")).collect(Collectors.toList());
        } catch (IOException e) { }
    }

    /**
     * download and unzip the cab file for this update
     */
    private void downloadCabFile() {
        try {
            // downloads cab file
            Files.copy(
                    new URL(this.downloadLink).openStream(),
                    Paths.get("drivers/".concat(uuid).concat(".cab")),
                    StandardCopyOption.REPLACE_EXISTING);

            // extract driver files
            new CabExtractor(
                    new File("drivers/".concat(uuid).concat(".cab")))
                    .extract();
        } catch (IOException | CabException e) { }
    }
}