Merge pull request #1348 from akto-api-security/feature/usage-cyborg

Feature/usage cyborg

apps/database-abstractor/src/main/java/com/akto/action/DbAction.java

@@ -42,6 +42,7 @@
 import com.akto.trafficFilter.HostFilter;
 import com.akto.trafficFilter.ParamFilter;
 import com.akto.util.Constants;
+import com.akto.dto.usage.MetricTypes;
 import com.akto.log.LoggerMaker;
 import com.akto.log.LoggerMaker.LogDb;
 import com.akto.util.enums.GlobalEnums.TestErrorSource;
@@ -224,6 +225,26 @@ public void setIssuesIds(BasicDBList issuesIds) {
     KafkaUtils kafkaUtils = new KafkaUtils();
     String endpointLogicalGroupId;
 
+    String metricType;
+
+    public String getMetricType() {
+        return metricType;
+    }
+
+    public void setMetricType(String metricType) {
+        this.metricType = metricType;
+    }
+
+    int deltaUsage;
+
+    public int getDeltaUsage() {
+        return deltaUsage;
+    }
+
+    public void setDeltaUsage(int deltaUsage) {
+        this.deltaUsage = deltaUsage;
+    }
+
     DataControlSettings dataControlSettings;
     public String fetchDataControlSettings() {
         try {
@@ -939,6 +960,27 @@ public String fetchStiBasedOnHostHeaders() {
         return Action.SUCCESS.toUpperCase();
     }
 
+    public String fetchDeactivatedCollections() {
+        try {
+            apiCollectionIds = DbLayer.fetchDeactivatedCollections();
+        } catch (Exception e) {
+            loggerMaker.errorAndAddToDb(e, "Error in fetchDeactivatedCollections " + e.toString());
+            return Action.ERROR.toUpperCase();
+        }
+        return Action.SUCCESS.toUpperCase();
+    }
+
+    public String updateUsage() {
+        try {
+            MetricTypes metric = MetricTypes.valueOf(metricType);
+            DbLayer.updateUsage(metric, deltaUsage);
+        } catch (Exception e) {
+            loggerMaker.errorAndAddToDb(e, "Error in updateUsage " + e.toString());
+            return Action.ERROR.toUpperCase();
+        }
+        return Action.SUCCESS.toUpperCase();
+    }
+
     public String fetchApiCollectionIds() {
         try {
             apiCollectionIds = DbLayer.fetchApiCollectionIds();

apps/database-abstractor/src/main/resources/struts.xml

@@ -233,6 +233,28 @@
             </result>
         </action>
 
+        <action name="api/fetchDeactivatedCollections" class="com.akto.action.DbAction" method="fetchDeactivatedCollections">
+            <interceptor-ref name="json"/>
+            <interceptor-ref name="defaultStack" />
+            <result name="SUCCESS" type="json"/>
+            <result name="ERROR" type="json">
+                <param name="statusCode">422</param>
+                <param name="ignoreHierarchy">false</param>
+                <param name="includeProperties">^actionErrors.*</param>
+            </result>
+        </action>
+
+        <action name="api/updateUsage" class="com.akto.action.DbAction" method="updateUsage">
+            <interceptor-ref name="json"/>
+            <interceptor-ref name="defaultStack" />
+            <result name="SUCCESS" type="json"/>
+            <result name="ERROR" type="json">
+                <param name="statusCode">422</param>
+                <param name="ignoreHierarchy">false</param>
+                <param name="includeProperties">^actionErrors.*</param>
+            </result>
+        </action>
+
         <action name="api/fetchEstimatedDocCount" class="com.akto.action.DbAction" method="fetchEstimatedDocCount">
             <interceptor-ref name="json"/>
             <interceptor-ref name="defaultStack" />

libs/dao/src/main/java/com/akto/dao/SetupDao.java

@@ -1,6 +1,7 @@
 package com.akto.dao;
 
 import com.akto.dto.Setup;
+import com.mongodb.BasicDBObject;
 
 public class SetupDao extends CommonContextDao<Setup> {
 
@@ -16,4 +17,17 @@ public Class<Setup> getClassT() {
         return Setup.class;
     }
 
+    private final static String SAAS = "saas";
+    public boolean isMetered() {
+        Setup setup = SetupDao.instance.findOne(new BasicDBObject());
+        boolean isSaas = false;
+        if (setup != null) {
+            String dashboardMode = setup.getDashboardMode();
+            if (dashboardMode != null) {
+                isSaas = dashboardMode.equalsIgnoreCase(SAAS);
+            }
+        }
+        return isSaas;
+    }
+
 }

libs/utils/src/main/java/com/akto/billing/UsageMetricUtils.java

@@ -13,6 +13,8 @@
 import org.json.JSONException;
 import org.json.JSONObject;
 
+import com.akto.RuntimeMode;
+import com.akto.dao.SetupDao;
 import com.akto.dao.billing.OrganizationsDao;
 import com.akto.dao.context.Context;
 import com.akto.dao.usage.UsageMetricInfoDao;
@@ -235,7 +237,13 @@ public static boolean checkTestRunsOverage(int accountId){
     public static FeatureAccess getFeatureAccess(int accountId, MetricTypes metricType) {
         FeatureAccess featureAccess = FeatureAccess.fullAccess;
         try {
-            if (!DashboardMode.isMetered()) {
+            /*
+             * runtime mode check for hybrid mini-runtime deployments.
+             * Setup Dao check for database-abstractor
+             */
+            if (!(DashboardMode.isMetered() ||
+                    RuntimeMode.isHybridDeployment() ||
+                    SetupDao.instance.isMetered())) {
                 return featureAccess;
             }
             Organization organization = OrganizationsDao.instance.findOneByAccountId(accountId);
@@ -249,7 +257,13 @@ public static FeatureAccess getFeatureAccess(int accountId, MetricTypes metricTy
     public static FeatureAccess getFeatureAccess(Organization organization, MetricTypes metricType) {
         FeatureAccess featureAccess = FeatureAccess.fullAccess;
         try {
-            if (!DashboardMode.isMetered()) {
+            /*
+             * runtime mode check for hybrid mini-runtime deployments.
+             * Setup Dao check for database-abstractor
+             */
+            if (!(DashboardMode.isMetered() ||
+                    RuntimeMode.isHybridDeployment() ||
+                    SetupDao.instance.isMetered())) {
                 return featureAccess;
             }
             if (organization == null) {

libs/utils/src/main/java/com/akto/data_actor/ClientActor.java

@@ -45,6 +45,7 @@
 import com.akto.dto.traffic.SampleData;
 import com.akto.dto.type.SingleTypeInfo;
 import com.akto.dto.type.URLMethods;
+import com.akto.dto.usage.MetricTypes;
 import com.akto.log.LoggerMaker;
 import com.akto.log.LoggerMaker.LogDb;
 import com.fasterxml.jackson.databind.DeserializationFeature;
@@ -297,6 +298,53 @@ public void bulkWriteTestingRunIssues(List<Object> writesForTestingRunIssues) {
         bulkWrite(writesForTestingRunIssues, "/bulkWriteTestingRunIssues", "writesForTestingRunIssues");
     }
 
+    public List<Integer> fetchDeactivatedCollections(){
+        List<Integer> ids = new ArrayList<>();
+
+        Map<String, List<String>> headers = buildHeaders();
+        OriginalHttpRequest request = new OriginalHttpRequest(url + "/fetchDeactivatedCollections", "", "GET", null, headers, "");
+        try {
+            OriginalHttpResponse response = ApiExecutor.sendRequest(request, true, null, false, null);
+            String responsePayload = response.getBody();
+            if (response.getStatusCode() != 200 || responsePayload == null) {
+                loggerMaker.errorAndAddToDb("invalid response in fetchDeactivatedCollections", LoggerMaker.LogDb.RUNTIME);
+                return ids;
+            }
+            BasicDBObject payloadObj;
+            try {
+                payloadObj =  BasicDBObject.parse(responsePayload);
+                BasicDBList objList = (BasicDBList) payloadObj.get("apiCollectionIds");
+                for (Object obj: objList) {
+                    int obj2 = (Integer) obj;
+                    ids.add(obj2);
+                }
+            } catch(Exception e) {
+                loggerMaker.errorAndAddToDb("error extracting response in fetchDeactivatedCollections" + e, LoggerMaker.LogDb.RUNTIME);
+            }
+        } catch (Exception e) {
+            loggerMaker.errorAndAddToDb("error in fetchDeactivatedCollections" + e, LoggerMaker.LogDb.RUNTIME);
+        }
+        return ids;
+    };
+
+    public void updateUsage(MetricTypes metricType, int deltaUsage){
+        Map<String, List<String>> headers = buildHeaders();
+        BasicDBObject obj = new BasicDBObject();
+        obj.put("metricType", metricType.name());
+        obj.put("deltaUsage", deltaUsage);
+        OriginalHttpRequest request = new OriginalHttpRequest(url + "/updateUsage", "", "POST", obj.toString(), headers, "");
+        try {
+            OriginalHttpResponse response = ApiExecutor.sendRequest(request, true, null, false, null);
+            String responsePayload = response.getBody();
+            if (response.getStatusCode() != 200 || responsePayload == null) {
+                loggerMaker.errorAndAddToDb("invalid response in updateUsage", LoggerMaker.LogDb.RUNTIME);
+            }
+        } catch (Exception e) {
+            loggerMaker.errorAndAddToDb("error in updateUsage" + e, LoggerMaker.LogDb.RUNTIME);
+        }
+        return;
+    };
+
     public TestSourceConfig findTestSourceConfig(String subType) {
         Map<String, List<String>> headers = buildHeaders();
         BasicDBObject obj = new BasicDBObject();

libs/utils/src/main/java/com/akto/data_actor/DataActor.java

@@ -20,10 +20,9 @@
 import com.akto.dto.testing.WorkflowTestResult;
 import com.akto.dto.testing.sources.TestSourceConfig;
 import com.akto.dto.traffic.SampleData;
-import com.akto.dto.traffic.TrafficInfo;
 import com.akto.dto.type.SingleTypeInfo;
 import com.akto.dto.type.URLMethods;
-import com.mongodb.client.model.WriteModel;
+import com.akto.dto.usage.MetricTypes;
 
 import java.util.List;
 import java.util.Map;
@@ -186,7 +185,9 @@ public abstract class DataActor {
 
     public abstract List<TestingRunIssues> fetchIssuesByIds(Set<TestingIssuesId> issuesIds);
 
-    // public abstract List<Testing
+    public abstract List<Integer> fetchDeactivatedCollections();
+
+    public abstract void updateUsage(MetricTypes metricType, int deltaUsage);
 
     public abstract List<SingleTypeInfo> findStiByParam(int apiCollectionId, String param);
 

libs/utils/src/main/java/com/akto/data_actor/DbActor.java

@@ -26,6 +26,7 @@
 import com.akto.dto.type.SingleTypeInfo;
 import com.akto.dto.type.URLMethods;
 import com.akto.dto.type.URLMethods.Method;
+import com.akto.dto.usage.MetricTypes;
 import com.mongodb.client.model.WriteModel;
 
 import java.util.ArrayList;
@@ -405,6 +406,16 @@ public TestingRunResultSummary updateIssueCountInSummary(String summaryId,
         return DbLayer.updateIssueCountInSummary(summaryId, totalCountIssues);
     }
 
+    public List<Integer> fetchDeactivatedCollections() {
+        return DbLayer.fetchDeactivatedCollections();
+    }
+
+    public void updateUsage(MetricTypes metricType,int deltaUsage){
+        DbLayer.updateUsage(metricType, deltaUsage);
+        return;
+    }
+
+
     public void updateIssueCountInTestSummary(String summaryId, Map<String, Integer> totalCountIssues) {
         DbLayer.updateIssueCountInTestSummary(summaryId, totalCountIssues, false);
     }

libs/utils/src/main/java/com/akto/data_actor/DbLayer.java

@@ -61,8 +61,11 @@
 import com.akto.dto.traffic_metrics.TrafficMetrics;
 import com.akto.dto.type.SingleTypeInfo;
 import com.akto.dto.type.URLMethods;
+import com.akto.dto.usage.MetricTypes;
 import com.akto.log.LoggerMaker;
 import com.akto.log.LoggerMaker.LogDb;
+import com.akto.usage.UsageMetricCalculator;
+import com.akto.usage.UsageMetricHandler;
 import com.akto.util.Constants;
 import com.mongodb.BasicDBObject;
 import com.mongodb.bulk.BulkWriteResult;
@@ -632,6 +635,16 @@ public static TestingRunResultSummary updateIssueCountInSummary(String summaryId
                 options);
     }
 
+    public static List<Integer> fetchDeactivatedCollections() {
+        return new ArrayList<>(UsageMetricCalculator.getDeactivatedLatest());
+    }
+
+    public static void updateUsage(MetricTypes metricType, int deltaUsage){
+        int accountId = Context.accountId.get();
+        UsageMetricHandler.calcAndFetchFeatureAccessUsingDeltaUsage(metricType, accountId, deltaUsage);
+        return;
+    }
+
     public static List<TestingRunResult> fetchLatestTestingRunResultBySummaryId(String summaryId, int limit, int skip) {
         ObjectId summaryObjectId = new ObjectId(summaryId);
         return TestingRunResultDao.instance

libs/utils/src/main/java/com/akto/usage/UsageMetricCalculator.java

@@ -23,9 +23,12 @@
 import com.akto.dto.usage.UsageMetric;
 import com.akto.dto.usage.metadata.ActiveAccounts;
 import com.akto.log.LoggerMaker;
+import com.akto.util.Constants;
 import com.akto.util.enums.GlobalEnums.YamlTemplateSource;
 import com.google.gson.Gson;
 import com.mongodb.client.model.Filters;
+import com.mongodb.client.model.Projections;
+
 import org.bson.conversions.Bson;
 
 public class UsageMetricCalculator {
@@ -60,7 +63,7 @@ public static Set<Integer> getDeactivated() {
 
     public static Set<Integer> getDeactivatedLatest(){
         List<ApiCollection> deactivated = ApiCollectionsDao.instance
-                .findAll(Filters.eq(ApiCollection._DEACTIVATED, true));
+                .findAll(Filters.eq(ApiCollection._DEACTIVATED, true), Projections.include(Constants.ID));
         Set<Integer> deactivatedIds = new HashSet<>(
                 deactivated.stream().map(apiCollection -> apiCollection.getId()).collect(Collectors.toList()));