Black formatting

plugins/velociraptor_legacy/icon_velociraptor_legacy/actions/run/action.py

@@ -8,23 +8,22 @@
 from pyvelociraptor import api_pb2_grpc
 import insightconnect_plugin_runtime
 from .schema import RunInput, RunOutput, Input, Output, Component
+
 # Custom imports below
 
 
 class Run(insightconnect_plugin_runtime.Action):
-
     def __init__(self):
         super(self.__class__, self).__init__(
-                name="run",
-                description=Component.DESCRIPTION,
-                input=RunInput(),
-                output=RunOutput())
+            name="run", description=Component.DESCRIPTION, input=RunInput(), output=RunOutput()
+        )
 
     def run(self, params={}):
         # START INPUT BINDING - DO NOT REMOVE - ANY INPUTS BELOW WILL UPDATE WITH YOUR PLUGIN SPEC AFTER REGENERATION
         # END INPUT BINDING - DO NOT REMOVE
         # TODO - If input bindings for connection can be done check to same if it you can do the same here
-        """Runs a VQL query against the Velociraptor server.
+        """
+        Runs a VQL query against the Velociraptor server.
 
         Args:
             config: A dictionary containing the configuration parameters for the Velociraptor server.
@@ -43,25 +42,32 @@ def run(self, params={}):
             certificate_chain_decoded = self.connection.certificate_chain_decoded
             query = params.get(Input.COMMAND)
             creds = grpc.ssl_channel_credentials(
-            root_certificates = root_certificates_decoded,
-            private_key = private_key_decoded,
-            certificate_chain = certificate_chain_decoded)
+                root_certificates=root_certificates_decoded,
+                private_key=private_key_decoded,
+                certificate_chain=certificate_chain_decoded,
+            )
             # This option is required to connect to the grpc server by IP - we
             # use self signed certs.
-            options = (('grpc.ssl_target_name_override', "VelociraptorServer",),)
+            options = (
+                (
+                    "grpc.ssl_target_name_override",
+                    "VelociraptorServer",
+                ),
+            )
             # The first step is to open a gRPC channel to the server..
-            with grpc.secure_channel(api_connection_string,
-                                    creds, options) as channel:
+            with grpc.secure_channel(api_connection_string, creds, options) as channel:
                 stub = api_pb2_grpc.APIStub(channel)
-                                # The request consists of one or more VQL queries. Note that you can collect artifacts by simply naming them using the
+                # The request consists of one or more VQL queries. Note that you can collect artifacts by simply naming them using the
                 # "Artifact" plugin.
                 request = api_pb2.VQLCollectorArgs(
                     max_wait=1,
                     max_row=100,
-                    Query=[api_pb2.VQLRequest(
-                        Name="ICON Plugin Request",
-                        VQL=query,
-                    )],
+                    Query=[
+                        api_pb2.VQLRequest(
+                            Name="ICON Plugin Request",
+                            VQL=query,
+                        )
+                    ],
                 )
                 # This will block as responses are streamed from the
                 # server. If the query is an event query we will run this loop
@@ -79,7 +85,6 @@ def run(self, params={}):
             results["logs_list"] = logs_list[0]
             return {Output.RESULTS: results}
         except grpc.RpcError as e:
-            self.logger.info("Error: ",e)
+            self.logger.info("Error: ", e)
             results["logs_list"] = e
             return {Output.RESULTS: results}
-

plugins/velociraptor_legacy/icon_velociraptor_legacy/connection/connection.py

@@ -8,21 +8,22 @@
 from pyvelociraptor import api_pb2_grpc
 import insightconnect_plugin_runtime
 from .schema import ConnectionSchema, Input
+
 # Custom imports below
 
 
 class Connection(insightconnect_plugin_runtime.Connection):
-
     def __init__(self):
         super(self.__class__, self).__init__(input=ConnectionSchema())
 
-    def connect(self, params):
+    def connect(self):
         # START INPUT BINDING - DO NOT REMOVE - ANY INPUTS BELOW WILL UPDATE WITH YOUR PLUGIN SPEC AFTER REGENERATION
         # TODO: generate bound input variables for the user, to help handhold the user
         # TODO: ex. self.api_key = params.get(Input.API_KEY)
         # END INPUT BINDING - DO NOT REMOVE
         self.logger.info("Connect: Connecting...")
-        """Runs a VQL query against the Velociraptor server.
+        """
+        Runs a VQL query against the Velociraptor server.
 
         Args:
             config: A dictionary containing the configuration parameters for the Velociraptor server.
@@ -41,15 +42,20 @@ def connect(self, params):
             certificate_chain_decoded = base64.b64decode(self.parameters["client_cert"]["secretKey"])
             query = "SELECT * FROM info()"
             creds = grpc.ssl_channel_credentials(
-            root_certificates = root_certificates_decoded,
-            private_key = private_key_decoded,
-            certificate_chain = certificate_chain_decoded)
+                root_certificates=root_certificates_decoded,
+                private_key=private_key_decoded,
+                certificate_chain=certificate_chain_decoded,
+            )
             # This option is required to connect to the grpc server by IP - we
             # use self signed certs.
-            options = (('grpc.ssl_target_name_override', "VelociraptorServer",),)
+            options = (
+                (
+                    "grpc.ssl_target_name_override",
+                    "VelociraptorServer",
+                ),
+            )
             # The first step is to open a gRPC channel to the server..
-            with grpc.secure_channel(api_connection_string,
-                                    creds, options) as channel:
+            with grpc.secure_channel(api_connection_string, creds, options) as channel:
                 stub = api_pb2_grpc.APIStub(channel)
 
                 # The request consists of one or more VQL queries. Note that
@@ -58,10 +64,12 @@ def connect(self, params):
                 request = api_pb2.VQLCollectorArgs(
                     max_wait=1,
                     max_row=100,
-                    Query=[api_pb2.VQLRequest(
-                        Name="ICON Plugin Request",
-                        VQL=query,
-                    )],
+                    Query=[
+                        api_pb2.VQLRequest(
+                            Name="ICON Plugin Request",
+                            VQL=query,
+                        )
+                    ],
                 )
                 # This will block as responses are streamed from the
                 # server. If the query is an event query we will run this loop
@@ -82,7 +90,7 @@ def connect(self, params):
             self.certificate_chain_decoded = base64.b64decode(self.parameters["client_cert"]["secretKey"])
             self.username = self.parameters["username"]
         except grpc.RpcError as e:
-            self.logger.info("Error: ",e)
+            self.logger.info("Error: ", e)
             self.api_connection_string = self.parameters["api_connection_string"]
             self.root_certificates_decoded = base64.b64decode(self.parameters["ca_certificate"]["secretKey"])
             self.private_key_decoded = base64.b64decode(self.parameters["client_private_key"]["secretKey"])

plugins/velociraptor_legacy/unit_test/test_run.py

@@ -1,6 +1,7 @@
 import sys
 import os
-sys.path.append(os.path.abspath('../'))
+
+sys.path.append(os.path.abspath("../"))
 
 from unittest import TestCase
 from icon_velociraptor_legacy.connection.connection import Connection