Task: generate an Load Balancer that forward to a NLB

Specification:
has one "aws_lb" resource
    with one of either "subnets" or "subnet_mapping"
    with load_balancer_type = "network"

has at least two aws_subnet
    has one of cidr_block, ipv6_cidr_block argument

has one aws_vpc
    has one of cider_block, ipv4_ipam_pool_id 

has one "aws_lb_listener" resource
    with one default_action argument
        with one type argument
            with one of following values: forward, redirect, fixed-response, authenticate-cognito and authenticate-oidc
    with one load_balancer_arn argument
    with one protocol argument
        with one of following values: TCP, UDP, TLS, TCP_UDP

has one "aws_lb_target_group"
    with port
    with protocol
    with vpc_id

has one aws_lb_target_group_attachment
    with one target_group_arn
    with one target_id

has at least one "aws_instance"
    with either launch_template or (ami and instance_type) argument

Terraform configuration (main.tf):
terraform {
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = ">= 5.75"
    }
  }

  required_version = ">= 1.9.0"
}

provider "aws" {
  region  = "us-east-1"
}

data "aws_availability_zones" "available" {
  state = "available"
}

module "vpc" {
  source  = "terraform-aws-modules/vpc/aws"

  name = "main-vpc"
  cidr = "10.0.0.0/16"

  azs                  = data.aws_availability_zones.available.names
  public_subnets       = ["10.0.4.0/24", "10.0.5.0/24", "10.0.6.0/24"]
}

data "aws_ami" "ubuntu" {
  most_recent = true
  owners      = ["amazon"]

  filter {
    name   = "name"
    values = ["*ubuntu-noble-24.04-amd64-server-*"]
  }
}

resource "aws_security_group" "instance-sg" {
  name = "instance-sg"
  vpc_id = module.vpc.vpc_id
}

resource "aws_vpc_security_group_ingress_rule" "instance-sg-ingress-rule" {
  from_port       = 80
  to_port         = 80
  ip_protocol     = "tcp"
  referenced_security_group_id = aws_security_group.lb-sg.id
  security_group_id = aws_security_group.instance-sg.id
}

resource "aws_vpc_security_group_egress_rule" "instance-sg-egress-rule" {
  from_port       = 0
  to_port         = 0
  ip_protocol     = "-1"
  cidr_ipv4       = "0.0.0.0/0"
  security_group_id = aws_security_group.instance-sg.id
}

resource "aws_security_group" "lb-sg" {
  name = "lb-sg"
  vpc_id = module.vpc.vpc_id
}

resource "aws_vpc_security_group_ingress_rule" "lb-sg-ingress-rule" {
  from_port   = 80
  to_port     = 80
  ip_protocol = "tcp"
  cidr_ipv4   = "0.0.0.0/0"
  security_group_id = aws_security_group.lb-sg.id
}

resource "aws_vpc_security_group_egress_rule" "lb-sg-egress-rule" {
  from_port   = 0
  to_port     = 0
  ip_protocol = "-1"
  cidr_ipv4   = "0.0.0.0/0"
  security_group_id = aws_security_group.lb-sg.id
}


resource "aws_instance" "my_instance" {
  ami           = data.aws_ami.ubuntu.id
  instance_type = "t2.micro"
  subnet_id     = module.vpc.public_subnets[0]
  security_groups = [ aws_security_group.instance-sg.id ]
}

resource "aws_lb" "test" {
  subnets            = module.vpc.public_subnets
  internal           = false
  load_balancer_type = "network"
  security_groups    = [aws_security_group.lb-sg.id]
}

resource "aws_lb_listener" "my_listener" {
  load_balancer_arn = aws_lb.test.arn
  port              = "80"
  protocol          = "TCP"

  default_action {
    type             = "forward"
    target_group_arn = aws_lb_target_group.target_group.arn
  }
}

resource "aws_lb_target_group" "target_group" {
  port     = 80
  protocol = "TCP"
  vpc_id   = module.vpc.vpc_id
}

resource "aws_lb_target_group_attachment" "ec2_attach" {
  target_group_arn = aws_lb_target_group.target_group.arn
  target_id = aws_instance.my_instance.id
}

Validation policy (checks.rego):
package terraform.validation

default is_valid_configuration = false

# Validate at least one aws_instance with the required arguments
is_valid_instance {
        count(valid_instances) > 0
}

valid_instances[instance] {
        instance := input.configuration.root_module.resources[_]
        instance.type == "aws_instance"
        requited_argument(instance)
}

requited_argument(instance) {
        instance.expressions.launch_template != null
}

requited_argument(instance) {
        instance.expressions.ami != null
        instance.expressions.instance_type != null
}

# Validate aws_lb
is_valid_lb {
        resource := input.configuration.root_module.resources[_]
        resource.type == "aws_lb"
        resource.expressions.load_balancer_type.constant_value == "network"
}

# Validate aws_lb_listener with the required arguments
is_valid_lb_listener {
        resource := input.configuration.root_module.resources[_]
        resource.type == "aws_lb_listener"
        resource.expressions.default_action
        resource.expressions.load_balancer_arn

        valid_protocols := {"TCP", "TLS", "UDP", "TCP_UDP"}
        protocol := resource.expressions.protocol.constant_value

        valid_protocols[protocol]
}

# Validate aws_lb_target_group exists
is_valid_lb_target_group {
        resource := input.configuration.root_module.resources[_]
        resource.type == "aws_lb_target_group"
}

# Validate aws_lb_target_group_attachment with the required arguments
is_valid_lb_target_group_attachment {
        resource := input.configuration.root_module.resources[_]
        resource.type == "aws_lb_target_group_attachment"
        resource.expressions.target_group_arn
        resource.expressions.target_id
}


# Aggregate validation
is_valid_configuration {
        is_valid_lb
        is_valid_lb_listener
        is_valid_lb_target_group_attachment
        is_valid_lb_target_group
        is_valid_instance
}