#!/usr/bin/env python3
"""
Helper script to produce NetworkPolicy that only allows DNS to all IPs and 
HTTPS to certain IPs.
Each IP should be provided on its own line to stdin. The NetworkPolicy yaml 
definition is provided to stdout.

Usage:
$ ./generate-deny-egress-yaml << EOF
127.0.1.2
169.23.23.23
EOF
"""

import sys

base = """
apiVersion: networking.k8s.io/v1
kind: NetworkPolicy
metadata:
  name: deny-egress-for-private-link-test
  namespace: sagemaker-k8s-operator-system
spec:
  podSelector: {}
  policyTypes:
  - Egress
  egress:
    # Allow DNS traffic
    - to:
      - ipBlock:
          cidr: 0.0.0.0/0
      ports:
      - protocol: UDP
        port: 53"""

per_ip = """
    # Allow HTTPS traffic to the specified IP
    - to:
      - ipBlock:
          cidr: %s/32
      ports:
        - protocol: TCP
          port: 443"""

print(base, end='')
for line in sys.stdin:
    print(per_ip % line.strip(), end='')
print('')