Set up prerequisites for python plugin support in Tsunami

PiperOrigin-RevId: 634773845
Change-Id: I2b0a4dacf1422403f3dbb8b5d9a412b957e47156

plugin_server/py/common/data/network_endpoint_utils.py

@@ -7,7 +7,7 @@
 
 import ipaddress
 from typing import Optional
-from tsunami.proto import network_pb2
+import network_pb2
 
 AddressFamily = network_pb2.AddressFamily
 Hostname = network_pb2.Hostname

plugin_server/py/common/data/network_endpoint_utils_test.py

@@ -2,8 +2,8 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-from tsunami.plugin_server.py.common.data import network_endpoint_utils
-from tsunami.proto import network_pb2
+from common.data import network_endpoint_utils
+import network_pb2
 
 AddressFamily = network_pb2.AddressFamily
 Hostname = network_pb2.Hostname

plugin_server/py/common/data/network_service_utils.py

@@ -8,9 +8,9 @@
 
 import socket
 import urllib.parse
-from tsunami.plugin_server.py.common.data import network_endpoint_utils
-from tsunami.proto import network_pb2
-from tsunami.proto import network_service_pb2
+from common.data import network_endpoint_utils
+import network_pb2
+import network_service_pb2
 
 urlparse = urllib.parse.urlparse
 NetworkService = network_service_pb2.NetworkService

plugin_server/py/common/data/network_service_utils_test.py

@@ -6,9 +6,9 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-from tsunami.plugin_server.py.common.data import network_service_utils
-from tsunami.proto import network_pb2
-from tsunami.proto import network_service_pb2
+from common.data import network_service_utils
+import network_pb2
+import network_service_pb2
 
 AddressFamily = network_pb2.AddressFamily
 NetworkEndpoint = network_pb2.NetworkEndpoint

plugin_server/py/common/net/http/host_resolver_http_adapter.py

@@ -6,7 +6,7 @@
 
 import requests
 
-from tsunami.plugin_server.py.common.net.http.http_header_fields import HttpHeaderFields
+from common.net.http.http_header_fields import HttpHeaderFields
 
 
 class HostResolverHttpAdapter(requests.adapters.HTTPAdapter):

plugin_server/py/common/net/http/host_resolver_http_adapter_test.py

@@ -6,9 +6,9 @@
 from absl.testing import parameterized
 import requests
 
-from tsunami.plugin_server.py.common.net.http.host_resolver_http_adapter import HostResolverHttpAdapter
-from tsunami.plugin_server.py.common.net.http.http_header_fields import HttpHeaderFields
-from tsunami.plugin_server.py.common.net.http.http_method import HttpMethod
+from common.net.http.host_resolver_http_adapter import HostResolverHttpAdapter
+from common.net.http.http_header_fields import HttpHeaderFields
+from common.net.http.http_method import HttpMethod
 
 
 class HostResolverHttpAdapterTest(parameterized.TestCase):

plugin_server/py/common/net/http/http_client.py

@@ -2,9 +2,9 @@
 import abc
 
 from typing import Optional
-from tsunami.plugin_server.py.common.data.network_service_utils import NetworkService
-from tsunami.plugin_server.py.common.net.http.http_request import HttpRequest
-from tsunami.plugin_server.py.common.net.http.http_response import HttpResponse
+from common.data.network_service_utils import NetworkService
+from common.net.http.http_request import HttpRequest
+from common.net.http.http_response import HttpResponse
 
 
 class HttpClient(metaclass=abc.ABCMeta):

plugin_server/py/common/net/http/http_header_fields_test.py

@@ -1,7 +1,7 @@
 """Tests for google3.third_party.java_src.tsunami.plugin_server.py.common.net.http_headers."""
 
 from absl.testing import absltest
-from tsunami.plugin_server.py.common.net.http import http_header_fields
+from common.net.http import http_header_fields
 
 
 HttpHeaderFields = http_header_fields.HttpHeaderFields

plugin_server/py/common/net/http/http_headers.py

@@ -3,7 +3,7 @@
 import collections
 import re
 from typing import Optional
-from tsunami.plugin_server.py.common.net.http.http_header_fields import HttpHeaderFields
+from common.net.http.http_header_fields import HttpHeaderFields
 
 
 class HttpHeaders:

plugin_server/py/common/net/http/http_headers_test.py

@@ -3,8 +3,8 @@
 import collections
 from absl.testing import absltest
 from absl.testing import parameterized
-from tsunami.plugin_server.py.common.net.http import http_header_fields
-from tsunami.plugin_server.py.common.net.http import http_headers
+from common.net.http import http_header_fields
+from common.net.http import http_headers
 
 
 HttpHeaderFields = http_header_fields.HttpHeaderFields

plugin_server/py/common/net/http/http_request.py

@@ -2,9 +2,9 @@
 
 from typing import Optional
 
-from tsunami.plugin_server.py.common.net.http.http_headers import Builder as HttpHeadersBuilder
-from tsunami.plugin_server.py.common.net.http.http_headers import HttpHeaders
-from tsunami.plugin_server.py.common.net.http.http_method import HttpMethod
+from common.net.http.http_headers import Builder as HttpHeadersBuilder
+from common.net.http.http_headers import HttpHeaders
+from common.net.http.http_method import HttpMethod
 
 
 def check_url_argument(func):

plugin_server/py/common/net/http/http_request_test.py

@@ -3,9 +3,9 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 
-from tsunami.plugin_server.py.common.net.http.http_headers import HttpHeaders
-from tsunami.plugin_server.py.common.net.http.http_method import HttpMethod
-from tsunami.plugin_server.py.common.net.http.http_request import HttpRequest
+from common.net.http.http_headers import HttpHeaders
+from common.net.http.http_method import HttpMethod
+from common.net.http.http_request import HttpRequest
 
 
 class HttpRequestTest(parameterized.TestCase):

plugin_server/py/common/net/http/http_response.py

@@ -3,8 +3,8 @@
 import abc
 import json
 from typing import Any, Optional
-from tsunami.plugin_server.py.common.net.http.http_headers import HttpHeaders
-from tsunami.plugin_server.py.common.net.http.http_status import HttpStatus
+from common.net.http.http_headers import HttpHeaders
+from common.net.http.http_status import HttpStatus
 
 
 class HttpResponse(metaclass=abc.ABCMeta):

plugin_server/py/common/net/http/http_response_test.py

@@ -2,9 +2,9 @@
 
 from absl.testing import absltest
 
-from tsunami.plugin_server.py.common.net.http.http_headers import HttpHeaders
-from tsunami.plugin_server.py.common.net.http.http_response import HttpResponse
-from tsunami.plugin_server.py.common.net.http.http_status import HttpStatus
+from common.net.http.http_headers import HttpHeaders
+from common.net.http.http_response import HttpResponse
+from common.net.http.http_status import HttpStatus
 
 
 class HttpResponseTest(absltest.TestCase):

plugin_server/py/common/net/http/http_status_test.py

@@ -2,7 +2,7 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-from tsunami.plugin_server.py.common.net.http import http_status
+from common.net.http import http_status
 
 HttpStatus = http_status.HttpStatus
 

plugin_server/py/common/net/http/requests_http_client.py

@@ -8,15 +8,15 @@
 from absl import logging
 import requests
 
-from tsunami.plugin_server.py.common.data.network_service_utils import NetworkService
-from tsunami.plugin_server.py.common.net.http.host_resolver_http_adapter import HostResolverHttpAdapter
-from tsunami.plugin_server.py.common.net.http.http_client import Builder
-from tsunami.plugin_server.py.common.net.http.http_client import HttpClient
-from tsunami.plugin_server.py.common.net.http.http_header_fields import HttpHeaderFields
-from tsunami.plugin_server.py.common.net.http.http_headers import HttpHeaders
-from tsunami.plugin_server.py.common.net.http.http_request import HttpRequest
-from tsunami.plugin_server.py.common.net.http.http_response import HttpResponse
-from tsunami.plugin_server.py.common.net.http.http_status import HttpStatus
+from common.data.network_service_utils import NetworkService
+from common.net.http.host_resolver_http_adapter import HostResolverHttpAdapter
+from common.net.http.http_client import Builder
+from common.net.http.http_client import HttpClient
+from common.net.http.http_header_fields import HttpHeaderFields
+from common.net.http.http_headers import HttpHeaders
+from common.net.http.http_request import HttpRequest
+from common.net.http.http_response import HttpResponse
+from common.net.http.http_status import HttpStatus
 
 
 _DEFAULT_ALLOW_REDIRECT = True

plugin_server/py/common/net/http/requests_http_client_test.py

@@ -6,18 +6,18 @@
 import requests
 import requests_mock
 
-from tsunami.plugin_server.py.common.data import network_endpoint_utils
-from tsunami.plugin_server.py.common.net.http.host_resolver_http_adapter import HostResolverHttpAdapter
-from tsunami.plugin_server.py.common.net.http.http_header_fields import HttpHeaderFields
-from tsunami.plugin_server.py.common.net.http.http_headers import HttpHeaders
-from tsunami.plugin_server.py.common.net.http.http_method import HttpMethod
-from tsunami.plugin_server.py.common.net.http.http_request import HttpRequest
-from tsunami.plugin_server.py.common.net.http.http_response import HttpResponse
-from tsunami.plugin_server.py.common.net.http.http_status import HttpStatus
-from tsunami.plugin_server.py.common.net.http.requests_http_client import RequestsHttpClient
-from tsunami.plugin_server.py.common.net.http.requests_http_client import RequestsHttpClientBuilder
-from tsunami.proto import network_pb2
-from tsunami.proto import network_service_pb2
+from common.data import network_endpoint_utils
+from common.net.http.host_resolver_http_adapter import HostResolverHttpAdapter
+from common.net.http.http_header_fields import HttpHeaderFields
+from common.net.http.http_headers import HttpHeaders
+from common.net.http.http_method import HttpMethod
+from common.net.http.http_request import HttpRequest
+from common.net.http.http_response import HttpResponse
+from common.net.http.http_status import HttpStatus
+from common.net.http.requests_http_client import RequestsHttpClient
+from common.net.http.requests_http_client import RequestsHttpClientBuilder
+import network_pb2
+import network_service_pb2
 
 
 class RequestsHttpClientTest(absltest.TestCase):

plugin_server/py/plugin/payload/payload.py

@@ -4,8 +4,8 @@
 
 from absl import logging
 
-from tsunami.plugin_server.py.plugin.payload.validator import Validator
-from tsunami.proto import payload_generator_pb2
+from plugin.payload.validator import Validator
+import payload_generator_pb2
 
 
 class Payload:

plugin_server/py/plugin/payload/payload_generator.py

@@ -3,11 +3,11 @@
 import re
 from typing import Any, Callable, Optional
 
-from tsunami.plugin_server.py.plugin.payload.payload import Payload
-from tsunami.plugin_server.py.plugin.payload.payload_secret_generator import PayloadSecretGenerator
-from tsunami.plugin_server.py.plugin.payload.validator import Validator
-from tsunami.plugin_server.py.plugin.tcs_client import TcsClient
-from tsunami.proto import payload_generator_pb2 as pg
+from plugin.payload.payload import Payload
+from plugin.payload.payload_secret_generator import PayloadSecretGenerator
+from plugin.payload.validator import Validator
+from plugin.tcs_client import TcsClient
+import payload_generator_pb2 as pg
 
 
 class PayloadGenerator:

plugin_server/py/plugin/payload/payload_generator_test.py

@@ -7,17 +7,17 @@
 from absl.testing import parameterized
 import requests_mock
 
-from google.protobuf import wrappers_pb2
-from tsunami.plugin_server.py.common.net.http.requests_http_client import RequestsHttpClientBuilder
-from tsunami.plugin_server.py.plugin.payload.payload_generator import PayloadGenerator
-from tsunami.plugin_server.py.plugin.payload.payload_generator_test_helper import ANY_SSRF_CONFIG
-from tsunami.plugin_server.py.plugin.payload.payload_generator_test_helper import JAVA_REFLECTIVE_RCE_CONFIG
-from tsunami.plugin_server.py.plugin.payload.payload_generator_test_helper import LINUX_REFLECTIVE_RCE_CONFIG
-from tsunami.plugin_server.py.plugin.payload.payload_generator_test_helper import LINUX_UNSPECIFIED_CONFIG
-from tsunami.plugin_server.py.plugin.payload.payload_secret_generator import PayloadSecretGenerator
-from tsunami.plugin_server.py.plugin.payload.payload_utility import get_parsed_payload
-from tsunami.plugin_server.py.plugin.tcs_client import TcsClient
-from tsunami.proto import payload_generator_pb2 as pg
+from google3.google.protobuf import wrappers_pb2
+from common.net.http.requests_http_client import RequestsHttpClientBuilder
+from plugin.payload.payload_generator import PayloadGenerator
+from plugin.payload.payload_generator_test_helper import ANY_SSRF_CONFIG
+from plugin.payload.payload_generator_test_helper import JAVA_REFLECTIVE_RCE_CONFIG
+from plugin.payload.payload_generator_test_helper import LINUX_REFLECTIVE_RCE_CONFIG
+from plugin.payload.payload_generator_test_helper import LINUX_UNSPECIFIED_CONFIG
+from plugin.payload.payload_secret_generator import PayloadSecretGenerator
+from plugin.payload.payload_utility import get_parsed_payload
+from plugin.tcs_client import TcsClient
+import payload_generator_pb2 as pg
 
 
 _IP_ADDRESS = '127.0.0.1'

plugin_server/py/plugin/payload/payload_generator_test_helper.py

@@ -1,5 +1,5 @@
 """Payload Generator Test Helper."""
-from tsunami.proto import payload_generator_pb2 as pg
+import payload_generator_pb2 as pg
 
 
 LINUX_REFLECTIVE_RCE_CONFIG = pg.PayloadGeneratorConfig(

plugin_server/py/plugin/payload/payload_secret_generator_test.py

@@ -2,8 +2,8 @@
 
 import base64
 
-from testing.pybase import googletest
-from tsunami.plugin_server.py.plugin.payload.payload_secret_generator import PayloadSecretGenerator
+from google3.testing.pybase import googletest
+from plugin.payload.payload_secret_generator import PayloadSecretGenerator
 
 
 _SECRET_BYTE_SIZE = 8

plugin_server/py/plugin/payload/payload_test.py

@@ -4,9 +4,9 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 
-from tsunami.plugin_server.py.plugin.payload.payload import Payload
-from tsunami.plugin_server.py.plugin.payload.validator import Validator
-from tsunami.proto import payload_generator_pb2 as pg
+from plugin.payload.payload import Payload
+from plugin.payload.validator import Validator
+import payload_generator_pb2 as pg
 
 _CONFIG = pg.PayloadGeneratorConfig()
 _PAYLOAD_ATTRIBUTES = pg.PayloadAttributes()

plugin_server/py/plugin/payload/payload_utility.py

@@ -1,13 +1,19 @@
 """Utility service for retrieving and parsing payload."""
 
 from ruamel import yaml
-from net.proto2.python.public import json_format
-from pyglib import resources
-from tsunami.proto import payload_generator_pb2 as pg
-
+from google.protobuf import json_format
+# copybara:strip_begin(internal-blaze-run)
+from google3.pyglib import resources
+# copybara:strip_end_and_replace_begin
+# from pathlib import Path
+# copybara:replace_end
+import payload_generator_pb2 as pg
 
+# copybara:strip_begin(internal-blaze-run)
 _PATH = 'google3/third_party/java_src/tsunami/plugin/src/main/resources/com/google/tsunami/plugin/payload/payload_definitions.yaml'
-
+# copybara:strip_end_and_replace_begin
+# _PATH = '../../plugin/src/main/resources/com/google/tsunami/plugin/payload/payload_definitions.yaml'
+# copybara:replace_end
 
 def get_parsed_payload() -> list[pg.PayloadDefinition]:
   """Get payload from payload_definitions.yaml.
@@ -25,7 +31,11 @@ def get_parsed_payload() -> list[pg.PayloadDefinition]:
         - Payload that uses validation regex but does not specify the regex to
           be used.
   """
+  # copybara:strip_begin(internal-blaze-run)
   payload_str = str(resources.GetResource(_PATH, 'r'))
+  # copybara:strip_end_and_replace_begin
+  # payload_str = Path(_PATH).read_text()
+  # copybara:replace_end
   payload_dict = yaml.safe_load(payload_str)
   payload_library = json_format.ParseDict(payload_dict, pg.PayloadLibrary())
   return _validate_payloads([p for p in payload_library.payloads])

plugin_server/py/plugin/payload/payload_utility_test.py

@@ -2,9 +2,9 @@
 
 import unittest
 
-from pyglib import resources
-from testing.pybase import googletest
-from tsunami.plugin_server.py.plugin.payload import payload_utility
+from google3.pyglib import resources
+from google3.testing.pybase import googletest
+from plugin.payload import payload_utility
 
 
 class PayloadUtilityTest(googletest.TestCase):

plugin_server/py/plugin/tcs_client.py

@@ -2,13 +2,13 @@
 import hashlib
 from typing import Optional
 from absl import logging
-from net.proto2.python.public import json_format
-from tsunami.plugin_server.py.common.data import network_endpoint_utils
-from tsunami.plugin_server.py.common.net.http.http_client import HttpClient
-from tsunami.plugin_server.py.common.net.http.http_headers import HttpHeaders
-from tsunami.plugin_server.py.common.net.http.http_request import HttpRequest
-from tsunami.proto import network_pb2
-from tsunami_callbackserver.proto import polling_pb2
+from google.protobuf import json_format
+from common.data import network_endpoint_utils
+from common.net.http.http_client import HttpClient
+from common.net.http.http_headers import HttpHeaders
+from common.net.http.http_request import HttpRequest
+import network_pb2
+import polling_pb2
 
 
 class TcsClient:

plugin_server/py/plugin/tcs_client_test.py

@@ -3,8 +3,8 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import requests_mock
-from tsunami.plugin_server.py.common.net.http.requests_http_client import RequestsHttpClientBuilder
-from tsunami.plugin_server.py.plugin.tcs_client import TcsClient
+from common.net.http.requests_http_client import RequestsHttpClientBuilder
+from plugin.tcs_client import TcsClient
 
 
 SECRET = 'a3d9ed89deadbeef'