def test_create_malware_policy(sut_handle, ixia_handle, policy_name,
policy_components, shieldx_constants,
shieldx_logger, datadir, pytestconfig):
# Initialize
# DUT
malware_mgmt = MalwareMgmt(sut_handle)
# JSON Config Reader
config_reader = CCR()
malware_json = "malware.json"
# Malware Payload - read config file
resolved_input_json_file = str((datadir / malware_json).resolve())
malware_payloads = dict(
config_reader.read_json_config(resolved_input_json_file))
# Create a new Malware Policy - QTAD-15358
payload = malware_payloads["malware_payload"]
shieldx_logger.info("Malware Payload - Template: {}".format(payload))
payload["name"] = policy_name
payload["fileActions"]["alertingThreshold"] = policy_components["alerting"]
payload["fileActions"]["blockingThreshold"] = policy_components["blocking"]
payload["fileActions"]["syslogThreshold"] = policy_components["syslog"]
payload["fileActions"]["selectedMWEngine"] = policy_components[
"malware_engine"]
shieldx_logger.info("Malware Payload - Create: {}".format(payload))
policy_id = malware_mgmt.create_policy(payload)
shieldx_logger.info("Policy ID: {}".format(policy_id))
def test_bats_001_create_resource_group(sut_handle, config_file, datadir,
shieldx_logger):
# Initialize
cloud_mgmt = CloudMgmt(sut_handle)
# JSON Config Reader
config_reader = CCR()
# Selected Rule IDs
resolved_input_json_file = str((datadir / config_file).resolve())
rg_config = config_reader.read_json_config(resolved_input_json_file)
# Create resource Group
resource_group = rg_config["rg1"]
resource_group["name"] = "VP_ResourceGroup"
resource_group["description"] = "CIDR Based RG"
resource_group["purpose"] = "POLICY"
resource_group["resourceType"] = "CIDR"
resource_group["memberList"] = [{
"id": 0,
"cidr": "192.168.131.5/32"
}, {
"id": 0,
"cidr": "192.168.131.51/32"
}]
rg_id = cloud_mgmt.create_resource_group(resource_group)
shieldx_logger.info("Resource Group Created, ID: {}".format(rg_id))
def test_delete_vpc(ec2_client, vpc_json, datadir, shieldx_constants,
shieldx_logger):
vpc = VPC(ec2_client=ec2_client)
config_reader = CCR()
# Read VPC Info
vpc_filename = str((datadir / vpc_json).resolve())
vpc_info = dict(config_reader.read_json_config(vpc_filename))
# Describe VPCs
vpc_cidr_block = vpc_info["cidr_block"]
vpc_response = vpc.describe_vpc_by_filter("cidr", vpc_cidr_block)
shieldx_logger.info("Describe VPC: {}".format(vpc_response))
# Fetch target VPC
target_vpc = get_target("CidrBlock", vpc_cidr_block, vpc_response["Vpcs"])
# Work with existing VPC
if target_vpc["found"]:
# Target VPC
vpc_id = target_vpc["target"]["VpcId"]
shieldx_logger.info("Target VPC ID: {}".format(vpc_id))
# Delete VPC
vpc_response = vpc.delete_vpc(vpc_id)
shieldx_logger.info("Delete VPC: {}".format(target_vpc["target"]))
else:
# NOOP - target VPC not found
pass
def test_bats_003_del_acl_rule(sut_handle, datadir, config_file,
shieldx_logger):
# Initialize
acl_mgmt = ACL(sut_handle)
# JSON Config Reader
config_reader = CCR()
# Selected Rule IDs
resolved_input_json_file = str((datadir / config_file).resolve())
acl_config = config_reader.read_json_config(resolved_input_json_file)
# Get Default Access Control Policy
default_access_control_policy = acl_mgmt.get_acl_policies()[0]
shieldx_logger.info(
"Before Del - Default ACP: {}".format(default_access_control_policy))
# Delete ACL Rule
index = 0
for acl_rule in default_access_control_policy["aclRules"]:
if acl_rule["name"] == acl_config["acl_rule1"]["name"]:
# Pop based on index
_ = default_access_control_policy["aclRules"].pop(index)
break
index += 1
shieldx_logger.info(
"After Del - Default ACP: {}".format(default_access_control_policy))
is_updated = acl_mgmt.update_acl(default_access_control_policy)
shieldx_logger.info("ACL Update status: {}".format(is_updated))
def test_enable_fireeye_malware_engine(sut_handle, shieldx_constants,
shieldx_logger, datadir, pytestconfig):
# Initialize
# DUT
malware_mgmt = MalwareMgmt(sut_handle)
#### Malware Engine API Key
feye_api_key = pytestconfig.getoption("malware_key")
assert feye_api_key is not None, "Malware Engine API Key is required."
# JSON Config Reader
config_reader = CCR()
malware_json = "malware.json"
# Malware Payload - read config file
resolved_input_json_file = str((datadir / malware_json).resolve())
malware_payloads = dict(
config_reader.read_json_config(resolved_input_json_file))
# Malware Engine payload
payload = malware_payloads["malware_engines_payload"]
shieldx_logger.info(
"Malware Engines Payload - Template: {}".format(payload))
payload["feCloudKey"] = feye_api_key
payload["feCloudEnabled"] = "true"
shieldx_logger.info("Malware Engines Payload - Enable: {}".format(payload))
is_updated = malware_mgmt.set_malware_config(payload)
shieldx_logger.info("Set Malware Config: {}".format(is_updated))
def test_disable_all_malware_engine(sut_handle, shieldx_constants,
shieldx_logger, datadir, pytestconfig):
# Initialize
# DUT
malware_mgmt = MalwareMgmt(sut_handle)
# JSON Config Reader
config_reader = CCR()
malware_json = "malware.json"
# Malware Payload - read config file
resolved_input_json_file = str((datadir / malware_json).resolve())
malware_payloads = dict(
config_reader.read_json_config(resolved_input_json_file))
# Malware Engine payload
payload = malware_payloads["malware_engines_payload"]
shieldx_logger.info(
"Malware Engines Payload - Template: {}".format(payload))
payload["feCloudKey"] = "null"
payload["feCloudEnabled"] = "false"
payload["sxCloudEnabled"] = "false"
shieldx_logger.info(
"Malware Engines Payload - Disable: {}".format(payload))
is_updated = malware_mgmt.set_malware_config(payload)
shieldx_logger.info("Set Malware Config: {}".format(is_updated))
def test_create_aws_connector(sut_handle, shieldx_logger, datadir,
input_json_file):
# Initialize
cloud_mgmt = CloudMgmt(sut_handle)
config_reader = CCR()
converter = DSC()
resolved_input_json_file = str((datadir / input_json_file).resolve())
cloud_payload = dict(
config_reader.read_json_config(resolved_input_json_file))
aws_cloud = cloud_payload["aws_cloud"]
# Fetch ACL Policy, Default ACL Policy = 3
acl_id = 3
aws_cloud["aclPolicyId"] = acl_id
# Inline Inspection: Active | Passive
pseudo_tap = "false"
aws_cloud["inlinePassiveInspection"] = pseudo_tap
# Create Infra Connector
cloud_id = cloud_mgmt.create_cloud(aws_cloud)
shieldx_logger.info("Cloud Type: {}".format("AWS"))
shieldx_logger.info("Cloud ID: {}".format(cloud_id))
shieldx_logger.info("---")
# Initialize
jobs_mgmt = JobsApis(sut_handle)
jobs = jobs_mgmt.get_jobs()
# Get Latest Job
shieldx_logger.info("Jobs count: {}".format(len(jobs)))
job = jobs[0]
job_id = job["id"]
# Monitor job progress
is_completed = False
retry = 0
max_retry = 10
time.sleep(60)
while retry < max_retry:
job = jobs_mgmt.get_job_by_id(job_id)
shieldx_logger.info("Job {} - {} - {}".format(job["id"], job["state"],
job["status"]))
if job["state"] == "COMPLETED":
break
retry += 1
time.sleep(60)
def test_delete_aws_connector(sut_handle, shieldx_logger, datadir,
input_json_file):
# Initialize
cloud_mgmt = CloudMgmt(sut_handle)
config_reader = CCR()
resolved_input_json_file = str((datadir / input_json_file).resolve())
cloud_payload = dict(
config_reader.read_json_config(resolved_input_json_file))
aws_cloud = cloud_payload["aws_cloud"]
name = aws_cloud["name"]
cloud_infra = cloud_mgmt.get_cloud_infra_by_name(name)
cloud_id = cloud_infra["id"]
shieldx_logger.info("Cloud Name: {}".format(name))
shieldx_logger.info("Cloud ID: {}".format(cloud_id))
shieldx_logger.info("---")
# Delete
cloud_mgmt.delete_cloud(cloud_id)
# Initialize
jobs_mgmt = JobsApis(sut_handle)
jobs = jobs_mgmt.get_jobs()
# Get Latest Job
shieldx_logger.info("Jobs count: {}".format(len(jobs)))
job = jobs[0]
job_id = job["id"]
# Monitor job progress
is_completed = False
retry = 0
max_retry = 10
time.sleep(60)
while not is_completed and retry < max_retry:
job = jobs_mgmt.get_job_by_id(job_id)
if job["state"] == "COMPLETED":
is_completed = True
shieldx_logger.info("Job {} - {} - {}".format(job["id"], job["state"],
job["status"]))
retry += 1
time.sleep(60)
def test_create_private_subnet(ec2_client, vpc_json, datadir,
shieldx_constants, shieldx_logger):
vpc = VPC(ec2_client=ec2_client)
config_reader = CCR()
# Read VPC Info
vpc_filename = str((datadir / vpc_json).resolve())
vpc_info = dict(config_reader.read_json_config(vpc_filename))
# Describe VPCs
vpc_cidr_block = vpc_info["cidr_block"]
vpc_response = vpc.describe_vpc_by_filter("cidr", vpc_cidr_block)
shieldx_logger.info("Describe VPC: {}".format(vpc_response))
# Fetch target VPC
target_vpc = get_target("CidrBlock", vpc_cidr_block, vpc_response["Vpcs"])
# If target VPC is found, proceed creating private subnet
if target_vpc["found"]:
# Get VPC Name and ID
vpc_target = target_vpc["target"]
shieldx_logger.info("Found matching VPC: {}".format(vpc_target))
vpc_id = vpc_target["VpcId"]
vpc_name = vpc_info["vpc_name"]
shieldx_logger.info("VPC Name: {}".format(vpc_name))
shieldx_logger.info("VPC ID: {}".format(vpc_id))
for private_subnet in vpc_info["private_subnet"]:
# Create Private Subnet
create_subnet_response = vpc.create_subnet(
vpc_id, private_subnet["address"])
private_subnet["id"] = create_subnet_response["Subnet"]["SubnetId"]
# Name Private Subnet
vpc.create_tags(private_subnet["id"], private_subnet["name"])
shieldx_logger.info("Create Private Subnet Response: {}".format(
create_subnet_response))
shieldx_logger.info("Private Subnet: {}".format(private_subnet))
# Write VPC Info - for debugging
shieldx_logger.info("VPC Info: {}".format(vpc_info))
shieldx_logger.info("VPC location: {}".format(vpc_filename))
config_reader.write_json_config(vpc_info, vpc_filename)
else:
# VPC not found, create VPC first
shieldx_logger.warning("VPC not found.")
def test_check_resources(vmware_cred, datadir, shieldx_logger):
config_reader = CCR()
# Get VMware Credentials
resolved_input_json_file = str((datadir / vmware_cred).resolve())
vmware_credentials = dict(
config_reader.read_json_config(resolved_input_json_file))
vcenter_ip = vmware_credentials.get("vcenter_ip", None)
username = vmware_credentials.get("username", None)
password = vmware_credentials.get("password", None)
# Connect
infra_client = VmwareInfra()
infra_client.connect(vcenter_ip, username, password)
# All Virtual Machines
all_vms = infra_client.get_all_vm()
for vm in all_vms:
shieldx_logger.info("Vitual machine: {}".format(vm.name))
# All Datacenters
all_dcs = infra_client.get_all_dc()
for dc in all_dcs:
shieldx_logger.info("Datacenter: {}".format(dc.name))
# All Datastores
all_dss = infra_client.get_all_ds()
for ds in all_dss:
shieldx_logger.info("Datastore: {}".format(ds.name))
# All Clusters
all_clusters = infra_client.get_all_clusters()
for cluster in all_clusters:
shieldx_logger.info("Cluster: {}".format(cluster.name))
# All Hosts
all_hosts = infra_client.get_all_hs()
for host in all_hosts:
shieldx_logger.info("Host: {}".format(host.name))
# All Networks
all_networks = infra_client.get_all_networks()
for network in all_networks:
shieldx_logger.info("Network: {}".format(network.name))
# Disconnect
infra_client.disconnect()
def test_tag_main_route_table(ec2_client, vpc_json, datadir, shieldx_constants,
shieldx_logger):
vpc = VPC(ec2_client=ec2_client)
config_reader = CCR()
# Read VPC Info
vpc_filename = str((datadir / vpc_json).resolve())
vpc_info = dict(config_reader.read_json_config(vpc_filename))
# Tags
vpc_owner = vpc_info["vpc_owner"]
rt_name = vpc_info["main_route_table"]["name"]
# Describe VPCs
vpc_cidr_block = vpc_info["cidr_block"]
vpc_response = vpc.describe_vpc_by_filter("cidr", vpc_cidr_block)
shieldx_logger.info("Describe VPC: {}".format(vpc_response))
# Fetch target VPC
target_vpc = get_target("CidrBlock", vpc_cidr_block, vpc_response["Vpcs"])
# Work with existing VPC
if target_vpc["found"]:
# Target VPC
vpc_id = target_vpc["target"]["VpcId"]
shieldx_logger.info("Target VPC ID: {}".format(vpc_id))
route_table_response = vpc.describe_route_tables_by_filter(
"vpc-id", vpc_id)
shieldx_logger.info("Describe RT: {}".format(route_table_response))
route_tables = list(route_table_response["RouteTables"])
# Main Route Table
main_rt = route_tables[0]
shieldx_logger.info("Main Route Table: {}".format(main_rt))
main_rt_id = main_rt["RouteTableId"]
shieldx_logger.info("Main Route Table ID: {}".format(main_rt_id))
# Tag main route table
tag_response = vpc.create_tags(main_rt_id, "Name", rt_name)
shieldx_logger.info("Tag Main Route Table: {}".format(tag_response))
tag_response = vpc.create_tags(main_rt_id, "Owner", vpc_owner)
shieldx_logger.info("Tag Main Route Table: {}".format(tag_response))
else:
# NOOP - target VPC not found
pass
def test_get_aws_cloud_objects(sut_handle, shieldx_logger, datadir,
input_json_file):
# Initialize
cloud_mgmt = CloudMgmt(sut_handle)
config_reader = CCR()
converter = DSC()
resolved_input_json_file = str((datadir / input_json_file).resolve())
cloud_payload = dict(
config_reader.read_json_config(resolved_input_json_file))
aws_cloud = cloud_payload["aws_cloud"]
name = aws_cloud["name"]
cloud_infra = cloud_mgmt.get_cloud_infra_by_name(name)
cloud_id = cloud_infra["id"]
shieldx_logger.info("Cloud Name: {}".format(name))
shieldx_logger.info("Cloud ID: {}".format(cloud_id))
shieldx_logger.info("---")
# Objects
cloud_objects = cloud_mgmt.get_cloud_objects(cloud_id)
#for key in cloud_objects:
# shieldx_logger.info("{}".format(key))
# #shieldx_logger.info("{} - {}".format(key, cloud_objects[key]))
networks = converter.list_of_dict_to_dict(cloud_objects["networks"],
"name")
shieldx_logger.info("Network: {} - {}".format(
networks["Juan-Management-Subnet"]["name"],
networks["Juan-Management-Subnet"]["id"]))
shieldx_logger.info("Network: {} - {}".format(
networks["Juan-Backplane-Subnet"]["name"],
networks["Juan-Management-Subnet"]["id"]))
shieldx_logger.info("Network: {} - {}".format(
networks["Juan-Workload-Subnet"]["name"],
networks["Juan-Management-Subnet"]["id"]))
tenants = converter.list_of_dict_to_dict(cloud_objects["tenants"], "name")
shieldx_logger.info("Tenant: {} - {} - {}".format(
tenants["Juan-Test-VPC"]["name"], tenants["Juan-Test-VPC"]["id"],
tenants["Juan-Test-VPC"]["regionId"]))
def test_modify_public_ec2_attribute(ec2_client, ec2_json, datadir,
shieldx_constants, shieldx_logger):
# Init clients
ec2 = EC2(ec2_client=ec2_client)
config_reader = CCR()
# Read EC2 Info
ec2_filename = str((datadir / ec2_json).resolve())
ec2_info = dict(config_reader.read_json_config(ec2_filename))
# Describe EC2 instances
ec2_response = dict(ec2.describe_ec2_instances())
reservations = list(ec2_response["Reservations"])
shieldx_logger.info("Reservations Count: {}".format(len(reservations)))
# Modify the public EC2 only
public_security_group = ec2_info["security_group"]["name"]
for reservation in reservations:
instance_id = reservation["Instances"][0]["InstanceId"]
state = reservation["Instances"][0]["State"]
security_groups = reservation["Instances"][0]["SecurityGroups"]
# Make the public EC2 instance immunte to API, or CLI termination
# It can only be terminated by going to AWS Console
# State - Code 16 = running
# State - Code 48 = terminated
if int(
state["Code"]
) == 16 and security_groups[0]["GroupName"] == public_security_group:
shieldx_logger.info(
"Modify this EC2 instance: {}".format(instance_id))
shieldx_logger.info(
"Modify this EC2 instance - State: {}".format(state))
shieldx_logger.info(
"Modify this EC2 Instance - Security Groups: {}".format(
security_groups))
ec2.modify_ec2_instance(instance_id)
else:
shieldx_logger.warning(
"Skipping this EC2 instance - ID: {}".format(instance_id))
shieldx_logger.warning(
"Skipping this EC2 instance - State: {}".format(state))
shieldx_logger.warning(
"Skipping this EC2 instance - Security Groups: {}".format(
security_groups))
def test_bats_002_add_acl_rule(sut_handle, datadir, config_file,
shieldx_logger):
# Initialize
acl_mgmt = ACL(sut_handle)
# JSON Config Reader
config_reader = CCR()
# Selected Rule IDs
resolved_input_json_file = str((datadir / config_file).resolve())
acl_config = config_reader.read_json_config(resolved_input_json_file)
# Get Default Access Control Policy
default_access_control_policy = acl_mgmt.get_acl_policies()[0]
shieldx_logger.info(
"Before Add - Default ACP: {}".format(default_access_control_policy))
# Clone ACL Rule and modify relevant fields
new_acl_rule = default_access_control_policy["aclRules"][0].copy()
del (new_acl_rule["id"])
new_acl_rule["name"] = acl_config["acl_rule1"]["name"]
new_acl_rule["description"] = acl_config["acl_rule1"]["description"]
# TODO
# new_acl_rule["spsId"] = (compute from vitual patch SPS ID)
# new_acl_rule["sourceResourceGroupList"] = (compute from RG or NS created based on WL IP from vuln scanner)
# new_acl_rule["destinationResourceGroupList"] = (compute from RG or NS created based on WL IP from vuln scanner)
# Append the new rule
default_access_control_policy["aclRules"].append(new_acl_rule)
# Fix order number, newly created rule is #1
acl_rules_count = len(default_access_control_policy["aclRules"])
for acl_rule in default_access_control_policy["aclRules"]:
acl_rule["orderNum"] = acl_rules_count
acl_rules_count -= 1
shieldx_logger.info(
"After Add - Default ACP: {}".format(default_access_control_policy))
is_updated = acl_mgmt.update_acl(default_access_control_policy)
shieldx_logger.info("ACL Update status: {}".format(is_updated))
def test_terminate_private_ec2(ec2_client, ec2_json, datadir,
shieldx_constants, shieldx_logger):
# Init clients
ec2 = EC2(ec2_client=ec2_client)
config_reader = CCR()
# Read EC2 Info
ec2_filename = str((datadir / ec2_json).resolve())
ec2_info = dict(config_reader.read_json_config(ec2_filename))
# Describe EC2 instances
ec2_response = dict(ec2.describe_ec2_instances())
reservations = list(ec2_response["Reservations"])
shieldx_logger.info("Reservations Count: {}".format(len(reservations)))
# Terminate the private EC2 instance
private_security_group = ec2_info["security_group"]["name"]
for reservation in reservations:
instance_id = reservation["Instances"][0]["InstanceId"]
state = reservation["Instances"][0]["State"]
security_groups = reservation["Instances"][0]["SecurityGroups"]
# Terminate the private EC2 instance
if int(
state["Code"]
) == 16 and security_groups[0]["GroupName"] == private_security_group:
shieldx_logger.info(
"Terminate this EC2 instance - ID: {}".format(instance_id))
shieldx_logger.info(
"Terminate this EC2 instance - State: {}".format(state))
shieldx_logger.info(
"Terminate this EC2 Instance - Security Groups: {}".format(
security_groups))
ec2.terminate_ec2_instance(instance_id)
else:
shieldx_logger.warning(
"Skipping this EC2 instance - ID: {}".format(instance_id))
shieldx_logger.warning(
"Skipping this EC2 instance - State: {}".format(state))
shieldx_logger.warning(
"Skipping this EC2 instance - Security Groups: {}".format(
security_groups))
def test_create_vswitch(vmware_cred, datadir, shieldx_logger):
config_reader = CCR()
# Get VMware Credentials
resolved_input_json_file = str((datadir / vmware_cred).resolve())
vmware_credentials = dict(
config_reader.read_json_config(resolved_input_json_file))
vcenter_ip = vmware_credentials.get("vcenter_ip", None)
username = vmware_credentials.get("username", None)
password = vmware_credentials.get("password", None)
esx_host = vmware_credentials.get("esx_host", None)
# Connect
infra_client = VmwareInfra()
infra_client.connect(vcenter_ip, username, password)
# All Hosts
all_hosts = infra_client.get_all_hs()
for host in all_hosts:
if host.name == esx_host:
shieldx_logger.info("Host: {}".format(host.name))
break
else:
pass
# vSwitch config
vswitch_config = {
"name": "Juan-TestCreate-vSwitch",
"num_ports": 120,
"mtu": 1500
}
# Create Standard vSwitch
infra_client.add_vswitch_to_host(host, vswitch_config)
# Add verification
# Disconnect
infra_client.disconnect()
def test_check_info(vmware_cred, datadir, shieldx_logger):
config_reader = CCR()
# Get VMware Credentials
resolved_input_json_file = str((datadir / vmware_cred).resolve())
vmware_credentials = dict(
config_reader.read_json_config(resolved_input_json_file))
vcenter_ip = vmware_credentials.get("vcenter_ip", None)
username = vmware_credentials.get("username", None)
password = vmware_credentials.get("password", None)
# Connect
infra_client = VmwareInfra()
infra_client.connect(vcenter_ip, username, password)
system_info = infra_client.get_system_info()
shieldx_logger.info("Name: {}".format(system_info["name"]))
shieldx_logger.info("Version: {}".format(system_info["version"]))
# Disconnect
infra_client.disconnect()
def test_check_networks(vmware_cred, datadir, shieldx_logger):
config_reader = CCR()
# Get VMware Credentials
resolved_input_json_file = str((datadir / vmware_cred).resolve())
vmware_credentials = dict(
config_reader.read_json_config(resolved_input_json_file))
vcenter_ip = vmware_credentials.get("vcenter_ip", None)
username = vmware_credentials.get("username", None)
password = vmware_credentials.get("password", None)
# Connect
infra_client = VmwareInfra()
infra_client.connect(vcenter_ip, username, password)
# All Networks
all_networks = infra_client.get_all_networks()
for network in all_networks:
shieldx_logger.info("Network: {}".format(network.name))
# Disconnect
infra_client.disconnect()
def test_bats_002_enable_license(sut_handle, license_json, shieldx_constants,
datadir):
manage = SysMgmt(sut_handle)
config_reader = CCR()
# Get license info
resolved_input_json_file = str((datadir / license_json).resolve())
license_info = dict(
config_reader.read_json_config(resolved_input_json_file))
license_key = license_info.get("key", None)
license_type = license_info.get("type", None)
license_capacity = license_info.get("capacity", None)
# Set the license
is_license_set = manage.set_license(license_key)
# Allow passage to cleanup tables
time.sleep(shieldx_constants["USER_WAIT"])
# Check if license is set
assert is_license_set == True, "Failed to activate license."
# Get the license info
license_info = manage.get_license(license_key)
# Check expected to actual capacity
fetched_capacity = license_info["expected_capacity"]
assert license_capacity == fetched_capacity, "Capacity does not match."
# Check expected to actual license type
fetched_license_type = license_info["license_type"]
assert license_type == fetched_license_type, "License type does not match."
# Check if license is active
fetched_license_state = license_info["current_state"]
assert "active" == fetched_license_state, "License is not active."
def test_create_deployment(sut_handle, shieldx_logger, datadir,
input_json_file):
# Initialize
cloud_mgmt = CloudMgmt(sut_handle)
dp_mgmt = DPMgmt(sut_handle)
config_reader = CCR()
converter = DSC()
resolved_input_json_file = str((datadir / input_json_file).resolve())
aws_payload = dict(
config_reader.read_json_config(resolved_input_json_file))
aws_cloud = aws_payload["aws_cloud"]
aws_deployment = aws_payload["aws_deployment"]
cloud_name = aws_cloud["name"]
cloud_infra = cloud_mgmt.get_cloud_infra_by_name(cloud_name)
cloud_id = cloud_infra["id"]
# Objects
cloud_objects = cloud_mgmt.get_cloud_objects(cloud_id)
networks = converter.list_of_dict_to_dict(cloud_objects["networks"],
"name")
tenants = converter.list_of_dict_to_dict(cloud_objects["tenants"], "name")
#for key in cloud_objects:
# shieldx_logger.info("{}".format(key))
# #shieldx_logger.info("{} - {}".format(key, cloud_objects[key]))
# Fill in deployment payload
tenant_name = aws_deployment["deploymentSpecification"]["tenantNameStr"]
mgmt_network_name = aws_deployment["deploymentSpecification"][
"mgmtNetworkNameStr"]
bkpln_network_name = aws_deployment["deploymentSpecification"][
"backPlaneNetworkStr"]
aws_deployment["cloudId"] = cloud_id
aws_deployment["deploymentSpecification"]["cloudId"] = cloud_id
aws_deployment["deploymentSpecification"]["cloudid"] = cloud_id
aws_deployment["deploymentSpecification"]["regionId"] = tenants[
tenant_name]["regionId"]
aws_deployment["deploymentSpecification"]["tenantId"] = tenants[
tenant_name]["id"]
aws_deployment["deploymentSpecification"]["availabilityZoneId"] = networks[
mgmt_network_name]["availabilityZoneId"]
aws_deployment["deploymentSpecification"]["mgmtNetworkId"] = networks[
mgmt_network_name]["id"]
aws_deployment["deploymentSpecification"]["backPlaneNetworkId"] = networks[
bkpln_network_name]["id"]
shieldx_logger.info("DP - {}".format(aws_deployment))
dp_id = dp_mgmt.create_deployment(aws_deployment)
shieldx_logger.info("Deployment ID: {}".format(dp_id))
shieldx_logger.info("---")
# Initialize
jobs_mgmt = JobsApis(sut_handle)
jobs = jobs_mgmt.get_jobs()
# Get Latest Job
shieldx_logger.info("Jobs count: {}".format(len(jobs)))
job = jobs[0]
job_id = job["id"]
# Monitor job progress
is_completed = False
retry = 0
max_retry = 10
time.sleep(60)
while retry < max_retry:
job = jobs_mgmt.get_job_by_id(job_id)
shieldx_logger.info("Job {} - {} - {}".format(job["id"], job["state"],
job["status"]))
if job["state"] == "COMPLETED":
break
retry += 1
time.sleep(60)
def test_create_management_subnet(ec2_client, vpc_json, datadir,
shieldx_constants, shieldx_logger):
vpc = VPC(ec2_client=ec2_client)
config_reader = CCR()
# Read VPC Info
vpc_filename = str((datadir / vpc_json).resolve())
vpc_info = dict(config_reader.read_json_config(vpc_filename))
# Describe VPCs
vpc_cidr_block = vpc_info["cidr_block"]
vpc_response = vpc.describe_vpc_by_filter("cidr", vpc_cidr_block)
shieldx_logger.info("Describe VPC: {}".format(vpc_response))
# Fetch target VPC
target_vpc = get_target("CidrBlock", vpc_cidr_block, vpc_response["Vpcs"])
# If target VPC is found, proceed creating public subnet
if target_vpc["found"]:
# Get VPC Name and ID
vpc_target = target_vpc["target"]
shieldx_logger.info("Found matching VPC: {}".format(vpc_target))
vpc_id = vpc_target["VpcId"]
vpc_name = vpc_info["vpc_name"]
vpc_owner = vpc_info["vpc_owner"]
shieldx_logger.info("VPC Name: {}".format(vpc_name))
shieldx_logger.info("VPC ID: {}".format(vpc_id))
for public_subnet in vpc_info["management_subnet"]:
# Create Subnet
create_subnet_response = vpc.create_subnet(
vpc_id, public_subnet["address"])
public_subnet["id"] = create_subnet_response["Subnet"]["SubnetId"]
# Tag Subnet
vpc.create_tags(public_subnet["id"], "Name", public_subnet["name"])
vpc.create_tags(public_subnet["id"], "Owner", vpc_owner)
shieldx_logger.info("Create Public Subnet Response: {}".format(
create_subnet_response))
shieldx_logger.info("Public Subnet: {}".format(public_subnet))
# Create Route Table
create_route_table_response = vpc.create_public_route_table(vpc_id)
shieldx_logger.info("Create Public Route Table: {}".format(
create_route_table_response))
# Get Route Table ID
route_table_id = create_route_table_response["RouteTable"][
"RouteTableId"]
vpc_info["management_route_table"]["id"] = route_table_id
# Tag Route Table
route_table_name = vpc_info["management_route_table"]["name"]
vpc.create_tags(route_table_id, "Name", route_table_name)
vpc.create_tags(route_table_id, "Owner", vpc_owner)
# Get IGW ID
describe_igw_response = vpc.describe_igw_by_filter(
"attachment.vpc-id", vpc_id)
shieldx_logger.info(
"Describe IGW Response: {}".format(describe_igw_response))
igw = describe_igw_response["InternetGateways"][0]
igw_id = igw["InternetGatewayId"]
vpc_info["internet_gateway"]["id"] = igw_id
shieldx_logger.info("Route Table ID: {}".format(route_table_id))
shieldx_logger.info("Internet Gateway ID: {}".format(igw_id))
# Create Route to IGW
create_route_response = vpc.create_igw_route_to_public_route_table(
route_table_id, igw_id)
shieldx_logger.info(
"Create Route to IGW: {}".format(create_route_response))
for public_subnet in vpc_info["management_subnet"]:
# Associate Management Subnet with Route Table
associate_subnet_response = vpc.associate_subnet_with_route_table(
route_table_id, public_subnet["id"])
shieldx_logger.info(
"Associate Public subnet with Route Table: {}".format(
associate_subnet_response))
# Allow auto-assign public IP Address for subnet
vpc.allow_auto_assign_ip_addresses_for_subnet(public_subnet["id"])
# Write VPC Info - for debugging
shieldx_logger.info("VPC Info: {}".format(vpc_info))
shieldx_logger.info("VPC location: {}".format(vpc_filename))
config_reader.write_json_config(vpc_info, vpc_filename)
else:
# VPC not found, create VPC first
shieldx_logger.warning("VPC not found.")
def test_create_vpc(ec2_client, vpc_json, datadir, shieldx_constants,
shieldx_logger):
vpc = VPC(ec2_client=ec2_client)
config_reader = CCR()
# Read VPC Info
vpc_filename = str((datadir / vpc_json).resolve())
vpc_info = dict(config_reader.read_json_config(vpc_filename))
# Describe VPCs
vpc_cidr_block = vpc_info["cidr_block"]
vpc_response = vpc.describe_vpc_by_filter("cidr", vpc_cidr_block)
shieldx_logger.info("Describe VPC: {}".format(vpc_response))
# Fetch target VPC
target_vpc = get_target("CidrBlock", vpc_cidr_block, vpc_response["Vpcs"])
# Work with existing VPC
if not target_vpc["found"]:
# Create VPC
vpc_response = vpc.create_vpc(vpc_cidr_block)
shieldx_logger.info("Create VPC: {}".format(vpc_response))
# Fetch VPC ID from create VPC response
vpc_id = vpc_response["Vpc"]["VpcId"]
vpc_info["vpc_id"] = vpc_id
vpc_name = vpc_info["vpc_name"]
vpc_owner = vpc_info["vpc_owner"]
shieldx_logger.info("VPC Name: {}".format(vpc_name))
shieldx_logger.info("VPC ID: {}".format(vpc_id))
# Tag the VPC
tag_response = vpc.create_tags(vpc_id, "Name", vpc_name)
shieldx_logger.info("Tag VPC: {}".format(tag_response))
tag_response = vpc.create_tags(vpc_id, "Owner", vpc_owner)
shieldx_logger.info("Tag VPC: {}".format(tag_response))
# Create Internet gateway
igw_response = vpc.create_internet_gateway()
shieldx_logger.info("Create IGW: {}".format(igw_response))
igw_id = igw_response["InternetGateway"]["InternetGatewayId"]
igw_name = vpc_info["internet_gateway"]["name"]
vpc_info["internet_gateway"]["id"] = igw_id
shieldx_logger.info("IGW Name: {}".format(igw_name))
shieldx_logger.info("IGW ID: {}".format(igw_id))
# Tag the IGW
tag_response = vpc.create_tags(igw_id, "Name", igw_name)
shieldx_logger.info("Tag IGW: {}".format(tag_response))
tag_response = vpc.create_tags(igw_id, "Owner", vpc_owner)
shieldx_logger.info("Tag IGW: {}".format(tag_response))
# Attach the Internet gateway
attach_response = vpc.attach_igw_to_vpc(igw_id, vpc_id)
shieldx_logger.info("Attach: {}".format(attach_response))
# Write VPC Info - for debugging
shieldx_logger.info("VPC Info: {}".format(vpc_info))
shieldx_logger.info("VPC location: {}".format(vpc_filename))
config_reader.write_json_config(vpc_info, vpc_filename)
else:
# VPC with given CIDR block already exist, skip create
shieldx_logger.warning("VPC Already Exist: {}".format(
target_vpc["target"]))
def test_launch_public_ec2_instance(ec2_client, vpc_json, ec2_json, datadir,
shieldx_constants, shieldx_logger):
# Init clients
vpc = VPC(ec2_client=ec2_client)
ec2 = EC2(ec2_client=ec2_client)
config_reader = CCR()
# Read VPC Info
vpc_filename = str((datadir / vpc_json).resolve())
vpc_info = dict(config_reader.read_json_config(vpc_filename))
# Read EC2 Info
ec2_filename = str((datadir / ec2_json).resolve())
ec2_info = dict(config_reader.read_json_config(ec2_filename))
# Describe VPCs
vpc_cidr_block = vpc_info["cidr_block"]
vpc_response = vpc.describe_vpc_by_filter("cidr", vpc_cidr_block)
shieldx_logger.info("Describe VPC: {}".format(vpc_response))
# Fetch target VPC
target_vpc = get_target("CidrBlock", vpc_cidr_block, vpc_response["Vpcs"])
# If target VPC is found, proceed creating public ec2 launch
if target_vpc["found"]:
# Get VPC Name and ID
vpc_target = target_vpc["target"]
shieldx_logger.info("Found matching VPC: {}".format(vpc_target))
vpc_id = vpc_target["VpcId"]
vpc_name = vpc_info["vpc_name"]
shieldx_logger.info("VPC Name: {}".format(vpc_name))
shieldx_logger.info("VPC ID: {}".format(vpc_id))
# Check if key pair exist
key_pair_name = ec2_info["keypair_name"]
key_pair_response = ec2.describe_key_pair_by_filter(
"key-name", key_pair_name)
target_key_pair = get_target("KeyName", key_pair_name,
key_pair_response["KeyPairs"])
# If key pair is missing, create it first
if not target_key_pair["found"]:
key_pair_response = ec2.create_key_pair(key_pair_name)
shieldx_logger.info(
"Key Pair Created: {}".format(key_pair_response))
else:
# Key pair already exist, proceed and use it
pass
# Check if Security Group exist
public_security_group_id = None
public_security_group_name = ec2_info["security_group"]["name"]
public_security_group_desc = ec2_info["security_group"]["desc"]
security_group_response = ec2.describe_security_group_by_filter(
"group-name", public_security_group_name)
target_security_group = get_target(
"GroupName", public_security_group_name,
security_group_response["SecurityGroups"])
# If target security group is missing, create it first
if not target_security_group["found"]:
# Create a Security Group
public_sg_response = ec2.create_security_group(
public_security_group_name, public_security_group_desc, vpc_id)
public_security_group_id = public_sg_response["GroupId"]
ec2_info["security_group"]["id"] = public_security_group_id
shieldx_logger.info("Security Group: {}".format(
ec2_info["security_group"]))
# IP Permissions
ip_permissions = ec2_info["ip_permissions"]
# Add inbound rule to security group - take from ec2_instance1.json
ec2.add_inbound_rule_to_sg(public_security_group_id,
ip_permissions)
else:
# Security Group Already exist, proceed and use it
public_security_group_id = target_security_group["target"][
"GroupId"]
# Image ID and Instance Type
image_id = ec2_info["image_id"]
instance_type = ec2_info["instance_type"]
# User Data - Startup script for EC2 instance
startup_script = """
#!/bin/bash
yum update -y
amazon-linux-extras install -y lamp-mariadb10.2-php7.2 php7.2
yum install -y httpd mariadb-server
systemctl start httpd
systemctl enable httpd
usermod -a -G apache ec2-user
chown -R ec2-user:apache /var/www
chmod 2775 /var/www
echo "<html><body><h1>Hello <b>ShieldX</b></h1></body></html>" > /var/www/html/index.html
"""
# EC2 CIDR block
ec2_cidr_block = ec2_info["cidr_block"]
subnets_response = vpc.describe_subnet_by_filter(
"cidr-block", ec2_cidr_block)
target_subnet = get_target("CidrBlock", ec2_cidr_block,
subnets_response["Subnets"])
# If subnet exist, proceed and use it
if target_subnet["found"]:
# Subnet ID
subnet_id = target_subnet["target"]["SubnetId"]
# Launch EC2 instance
launch_ec2_response = ec2.launch_ec2_instance(
image_id, key_pair_name, 1, 1, instance_type,
public_security_group_id, subnet_id, startup_script)
shieldx_logger.info(
"Launching EC2 Instance: {}".format(launch_ec2_response))
else:
# Public Subnet not found
shieldx_logger.warning(
"Public subnet not found: {}".format(ec2_cidr_block))
else:
# VPC not found, create VPC first
shieldx_logger.warning("VPC not found.")
def test_launch_private_ec2_instance(ec2_client, vpc_json, ec2_json, datadir,
shieldx_constants, shieldx_logger):
# Init clients
vpc = VPC(ec2_client=ec2_client)
ec2 = EC2(ec2_client=ec2_client)
config_reader = CCR()
# Read VPC Info
vpc_filename = str((datadir / vpc_json).resolve())
vpc_info = dict(config_reader.read_json_config(vpc_filename))
# Read EC2 Info
ec2_filename = str((datadir / ec2_json).resolve())
ec2_info = dict(config_reader.read_json_config(ec2_filename))
# Describe VPCs
vpc_cidr_block = vpc_info["cidr_block"]
vpc_response = vpc.describe_vpc_by_filter("cidr", vpc_cidr_block)
shieldx_logger.info("Describe VPC: {}".format(vpc_response))
# Fetch target VPC
target_vpc = get_target("CidrBlock", vpc_cidr_block, vpc_response["Vpcs"])
# If target VPC is not present, create
if target_vpc["found"]:
# Get VPC Name and ID
vpc_target = target_vpc["target"]
shieldx_logger.info("Found matching VPC: {}".format(vpc_target))
vpc_id = vpc_target["VpcId"]
vpc_name = vpc_info["vpc_name"]
shieldx_logger.info("VPC Name: {}".format(vpc_name))
shieldx_logger.info("VPC ID: {}".format(vpc_id))
# Check if key pair exist
key_pair_name = ec2_info["keypair_name"]
key_pair_response = ec2.describe_key_pair_by_filter(
"key-name", key_pair_name)
target_key_pair = get_target("KeyName", key_pair_name,
key_pair_response["KeyPairs"])
# If key pair is missing, create it first
if not target_key_pair["found"]:
key_pair_response = ec2.create_key_pair(key_pair_name)
shieldx_logger.info(
"Key Pair Created: {}".format(key_pair_response))
else:
# Key pair already exist, proceed and use it
pass
# Check if Security Group exist
private_security_group_id = None
private_security_group_name = ec2_info["security_group"]["name"]
private_security_group_desc = ec2_info["security_group"]["desc"]
security_group_response = ec2.describe_security_group_by_filter(
"group-name", private_security_group_name)
target_security_group = get_target(
"GroupName", private_security_group_name,
security_group_response["SecurityGroups"])
# If target security group is missing, create it first
if not target_security_group["found"]:
# Create a Security Group
private_sg_response = ec2.create_security_group(
private_security_group_name, private_security_group_desc,
vpc_id)
private_security_group_id = private_sg_response["GroupId"]
ec2_info["security_group"]["id"] = private_security_group_id
shieldx_logger.info("Security Group: {}".format(
ec2_info["security_group"]))
# IP Permissions - no access from the internet
ip_permissions = ec2_info["ip_permissions"]
# Add inbound rule to security group - take from ec2_instance1.json
ec2.add_inbound_rule_to_sg(private_security_group_id,
ip_permissions)
else:
# Security Group Already exist, proceed and use it
private_security_group_id = target_security_group["target"][
"GroupId"]
# Image ID and Instance Type
image_id = ec2_info["image_id"]
instance_type = ec2_info["instance_type"]
# User Data - Startup script for EC2 instance
startup_script = """"""
# EC2 CIDR block
ec2_cidr_block = ec2_info["cidr_block"]
subnets_response = vpc.describe_subnet_by_filter(
"cidr-block", ec2_cidr_block)
target_subnet = get_target("CidrBlock", ec2_cidr_block,
subnets_response["Subnets"])
# If subnet exist, proceed and use it
if target_subnet["found"]:
# Subnet ID
subnet_id = target_subnet["target"]["SubnetId"]
# Launch EC2 instance
launch_ec2_response = ec2.launch_ec2_instance(
image_id, key_pair_name, 1, 1, instance_type,
private_security_group_id, subnet_id, startup_script)
shieldx_logger.info(
"Launching EC2 Instance: {}".format(launch_ec2_response))
else:
# Private Subnet not found
shieldx_logger.warning(
"Private subnet not found: {}".format(ec2_cidr_block))
else:
# VPC not found, create VPC first
shieldx_logger.warning("VPC not found.")
def test_bats_025_tpp_by_rule_ids(sut_handle, datadir, shieldx_logger,
input_json_file, tpp_name, block_threats):
""" Create TPP by rule IDs, block threats if flag is set. """
# JSON Config Reader
config_reader = CCR()
# Based on the All Threats TPP
policy_name = "All Threats"
from_policy_id = None
to_policy_id = None
threats = None
apps = None
policy_mgmt = PolicyMgmt(sut_handle)
# Get Threat Prevention Policy
tpp = policy_mgmt.get_threat_prevention_policy_by_name(policy_name)
from_policy_id = tpp["id"]
if from_policy_id is not None:
threats = policy_mgmt.get_threats_by_policy_id(from_policy_id)
apps = policy_mgmt.get_apps_by_policy_id(from_policy_id)
shieldx_logger.info("TPP Name: {}".format(policy_name))
shieldx_logger.info("TPP ID: {}".format(from_policy_id))
else:
shieldx_logger.error("Unable to find the TPP.")
# Selected Rule IDs
resolved_input_json_file = str((datadir / input_json_file).resolve())
selected_app_names = list(
config_reader.read_json_config(resolved_input_json_file))
shieldx_logger.info("Selected Names: {}".format(selected_app_names))
# Fetch the payload from a config file
tpp_config_file = "tpp.json"
resolved_tpp_config_file = str((datadir / tpp_config_file).resolve())
tpp_config = read_config(resolved_tpp_config_file)
if tpp_config is not None:
# Get the payload for cloning
if "tpp_clone_payload" in tpp_config:
clone_payload = tpp_config["tpp_clone_payload"]
shieldx_logger.info("Clone Payload: {}".format(clone_payload))
else:
pass
# Get the payload for response action
if "tpp_bulk_edit_response_payload" in tpp_config:
response_payload = tpp_config["tpp_bulk_edit_response_payload"]
shieldx_logger.info(
"Bulk Edit Payload: {}".format(response_payload))
else:
pass
# Populate the payload
if clone_payload:
clone_payload["name"] = tpp_name
clone_payload["tenantId"] = 1 # this should be fetched
clone_payload["rules"] = [{"appNames": selected_app_names}]
else:
shieldx_logger.error(
"Unable to fetch the TPP payload from config file.")
# Create a clone of a TPP, get the TPP ID back
to_policy_id = policy_mgmt.create_threat_prevention_policy(clone_payload)
shieldx_logger.info("Create OK, Policy ID: {}".format(to_policy_id))
assert to_policy_id != 0, "TPP Clone failed."
# Bulk Edit - Block threats
if block_threats:
threat_responses = policy_mgmt.get_threat_responses_by_policy_id(
to_policy_id)
if len(threat_responses) > 0:
for threat_response in threat_responses:
threat_response["block"] = True
threat_response["policyId"] = to_policy_id
response_payload["id"] = to_policy_id
response_payload["responses"] = threat_responses
shieldx_logger.info(
"Bulk Edit Payload: {}".format(response_payload))
bulk_edit_success = policy_mgmt.bulk_update_threat_responses(
response_payload)
assert bulk_edit_success == True, "Bulk edit response action failed."
else:
shieldx_logger.error("NOOP - no threat found.")
else:
shieldx_logger.info("NOOP - no blocking required.")