pf-sdn/app/src/main/java/ru/kirillius/pf/sdn/App.java

package ru.kirillius.pf.sdn;

import lombok.Getter;
import lombok.SneakyThrows;
import ru.kirillius.json.rpc.Servlet.JSONRPCServlet;
import ru.kirillius.pf.sdn.External.API.Components.FRR;
import ru.kirillius.pf.sdn.External.API.Components.OVPN;
import ru.kirillius.pf.sdn.External.API.Components.TDNS;
import ru.kirillius.pf.sdn.External.API.HEInfoProvider;
import ru.kirillius.pf.sdn.core.Auth.AuthManager;
import ru.kirillius.pf.sdn.core.Auth.TokenStorage;
import ru.kirillius.pf.sdn.core.*;
import ru.kirillius.pf.sdn.core.Networking.ASInfoService;
import ru.kirillius.pf.sdn.core.Networking.NetworkManager;
import ru.kirillius.pf.sdn.core.Subscription.SubscriptionManager;
import ru.kirillius.pf.sdn.core.Util.Wait;
import ru.kirillius.pf.sdn.web.HTTPServer;
import ru.kirillius.utils.logging.SystemLogger;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.logging.Level;

public class App implements Context, Closeable {
    private final static File configFile = new File("config.json");
    protected final static String CTX = App.class.getSimpleName();

    static {
        SystemLogger.initializeLogging(Level.INFO, List.of(InMemoryLogHandler.class));
        SystemLogger.setExceptionDumping(true);
    }

    private final AtomicBoolean shouldRestart = new AtomicBoolean(false);
    private final AtomicBoolean running = new AtomicBoolean(true);

    @Getter
    private final NetworkManager networkManager;
    @Getter
    private volatile Config config;
    @Getter
    private final AuthManager authManager;
    @Getter
    private final HTTPServer server;
    @Getter
    private final ASInfoService ASInfoService;
    @Getter
    private final SubscriptionManager subscriptionManager;
    @Getter
    private final UpdateManager updateManager;
    @Getter
    private final TokenStorage tokenStorage;
    @Getter
    private final ContextEventsHandler EventsHandler = new ContextEventsHandler();

    private final List<Component<?>> loadedComponents = new ArrayList<>();

    @SneakyThrows
    public App(File configFile) {
        try {
            config = Config.load(configFile);
        } catch (IOException e) {
            config = new Config();
            try {
                Config.store(config, configFile);
            } catch (IOException ex) {
                throw new RuntimeException(ex);
            }
        }

        authManager = new AuthManager(this);
        ASInfoService = new ASInfoService();
        ASInfoService.setProvider(new HEInfoProvider(this));
        networkManager = new NetworkManager(this);
        networkManager.getInputResources().add(config.getCustomResources());
        subscriptionManager = new SubscriptionManager(this);
        updateManager = new UpdateManager(this);
        tokenStorage = new TokenStorage(this);
        subscribe();
        updateManager.start();
        initComponents();
        server = new HTTPServer(this);

        if (config.getPasswordHash() == null || config.getPasswordHash().isEmpty()) {
            SystemLogger.error("There is no password for admin. Setting default password: admin", CTX);
            getAuthManager().updatePassword("admin");
        }
        getSubscriptionManager().triggerUpdate();
    }

    public static void main(String[] args) {
        var restart = false;
        do {
            try (var app = new App(configFile)) {
                Wait.when(app.running::get);
                restart = app.shouldRestart.get();
            } catch (Exception e) {
                SystemLogger.error("Unhandled error", CTX, e);
            }
        } while (restart);
    }

    public void triggerRestart() {
        SystemLogger.message("Restarting app", CTX);
        running.set(false);
        shouldRestart.set(true);
    }

    public void triggerShutdown() {
        SystemLogger.message("Shutting down app", CTX);
        running.set(false);
    }

    public Collection<Class<? extends Component<?>>> getComponentClasses() {
        return List.of(FRR.class, OVPN.class, TDNS.class);
    }

    private void unloadComponent(Component<?> component) {
        SystemLogger.message("Unloading component: " + component.getClass().getSimpleName(), CTX);
        try {
            component.close();
        } catch (IOException e) {
            SystemLogger.error("Error on component unload", CTX, e);
        } finally {
            loadedComponents.remove(component);
        }
    }

    private void loadComponent(Class<? extends Component<?>> componentClass) {
        SystemLogger.message("Loading component: " + componentClass.getSimpleName(), CTX);
        var plugin = Component.loadPlugin(componentClass, this);
        loadedComponents.add(plugin);
    }

    public void initComponents() {
        var enabledPlugins = config.getEnabledComponents();

        (List.copyOf(loadedComponents)).forEach(plugin -> {
            if (!enabledPlugins.contains(plugin.getClass())) {
                unloadComponent(plugin);
            }
        });
        var loadedClasses = loadedComponents.stream().map(plugin -> plugin.getClass()).toList();
        enabledPlugins.forEach(pluginClass -> {
            if (loadedClasses.contains(pluginClass)) {
                return;
            }
            loadComponent(pluginClass);
        });
    }

    @Override
    public void reloadComponents(Class<? extends Component<?>>... classes) {
        Arrays.stream(classes)
                .forEach(componentClass -> {
                            loadedComponents.stream()
                                    .filter(component -> componentClass.equals(component.getClass()))
                                    .findFirst().ifPresent(this::unloadComponent);
                            loadComponent(componentClass);
                        }
                );

    }

    @Override
    public JSONRPCServlet getRPC() {
        if (server == null) {
            return null;
        }
        return server.getJSONRPC();
    }

    private void subscribe() {
        var eventsHandler = getEventsHandler();
        eventsHandler.getSubscriptionsUpdateEvent().add(bundle -> {
            var manager = getNetworkManager();
            var inputResources = getNetworkManager().getInputResources();
            inputResources.clear();
            inputResources.add(config.getCustomResources());
            inputResources.add(bundle);
            manager.triggerUpdate(false);
        });
    }

    @Override
    public Component<?> getComponentInstance(Class<? extends Component<?>> pluginClass) {
        return loadedComponents.stream().filter(plugin -> plugin.getClass().equals(pluginClass)).findFirst().orElse(null);
    }

    @Override
    public void close() throws IOException {
        loadedComponents.forEach(plugin -> {
            try {
                plugin.close();
            } catch (IOException e) {
                SystemLogger.error("Error closing plugin", CTX, e);
            }
        });
        ASInfoService.close();
        networkManager.close();
        try {
            server.stop();
        } catch (Exception e) {
            SystemLogger.error("Error stopping server", CTX, e);
        }
    }
}