#!/usr/bin/python
#
# Copyright 2018-2021 Polyaxon, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest

from polyaxon import settings
from polyaxon.proxies.schemas.gateway.api import get_api_location_config
from polyaxon.proxies.schemas.gateway.auth import (
    get_auth_config,
    get_auth_location_config,
)
from polyaxon.proxies.schemas.gateway.dns import get_dns_config, get_resolver
from polyaxon.proxies.schemas.gateway.redirect import get_redirect_config
from polyaxon.proxies.schemas.gateway.services import (
    get_plugins_location_config,
    get_services_location_config,
)
from polyaxon.proxies.schemas.gateway.ssl import get_ssl_config
from polyaxon.proxies.schemas.gateway.streams import get_streams_location_config
from polyaxon.utils.test_utils import BaseTestCase


@pytest.mark.proxies_mark
class TestGatewaySchemas(BaseTestCase):
    SET_PROXIES_SETTINGS = True

    def test_ssl(self):
        expected = r"""
# SSL
ssl_session_timeout 1d;
ssl_session_cache shared:SSL:50m;
ssl_session_tickets off;

# intermediate configuration
ssl_protocols TLSv1.2 TLSv1.3;
ssl_ciphers ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:DHE-RSA-AES128-GCM-SHA256:DHE-RSA-AES256-GCM-SHA384;
ssl_prefer_server_ciphers on;

# OCSP Stapling
ssl_stapling on;
ssl_stapling_verify on;
resolver 1.1.1.1 1.0.0.1 [2606:4700:4700::1111] [2606:4700:4700::1001] 8.8.8.8 8.8.4.4 [2001:4860:4860::8888] [2001:4860:4860::8844] 208.67.222.222 208.67.220.220 [2620:119:35::35] [2620:119:53::53] valid=60s;
resolver_timeout 2s;

ssl_certificate      /etc/ssl/polyaxon/polyaxon.com.crt;
ssl_certificate_key  /etc/ssl/polyaxon/polyaxon.com.key;
"""  # noqa
        assert get_ssl_config() == expected

        expected = r"""
# SSL
ssl_session_timeout 1d;
ssl_session_cache shared:SSL:50m;
ssl_session_tickets off;

# intermediate configuration
ssl_protocols TLSv1.2 TLSv1.3;
ssl_ciphers ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:DHE-RSA-AES128-GCM-SHA256:DHE-RSA-AES256-GCM-SHA384;
ssl_prefer_server_ciphers on;

# OCSP Stapling
ssl_stapling on;
ssl_stapling_verify on;
resolver 1.1.1.1 1.0.0.1 [2606:4700:4700::1111] [2606:4700:4700::1001] 8.8.8.8 8.8.4.4 [2001:4860:4860::8888] [2001:4860:4860::8844] 208.67.222.222 208.67.220.220 [2620:119:35::35] [2620:119:53::53] valid=60s;
resolver_timeout 2s;

ssl_certificate      /foo/polyaxon.com.crt;
ssl_certificate_key  /foo/polyaxon.com.key;
"""  # noqa
        settings.PROXIES_CONFIG.ssl_path = "/foo"
        assert get_ssl_config() == expected

    def test_redirect_config(self):
        expected = r"""
server {
    listen 80;
    return 301 https://$host$request_uri;
}
"""  # noqa
        settings.PROXIES_CONFIG.ssl_enabled = False
        assert get_redirect_config() == ""
        settings.PROXIES_CONFIG.ssl_enabled = True
        assert get_redirect_config() == expected


@pytest.mark.proxies_mark
class TestGatewayServicesSchemas(BaseTestCase):
    SET_PROXIES_SETTINGS = True

    def test_service_dns_resolver(self):
        settings.PROXIES_CONFIG.auth_enabled = False
        expected = r"""
location ~ /services/v1/([-_.:\w]+)/([-_.:\w]+)/([-_.:\w]+)/runs/([-_.:\w]+)/(.*) {
    proxy_pass http://plx-operation-$4.$1.svc.cluster.local;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_hide_header X-Frame-Options;
    proxy_set_header Origin "";
    proxy_set_header Host $http_host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_buffering off;
}
"""  # noqa
        settings.PROXIES_CONFIG.dns_use_resolver = False
        resolver = get_resolver()
        assert (
            get_services_location_config(resolver=resolver, auth="", rewrite=False)
            == expected
        )

        expected = r"""
location ~ /services/v1/([-_.:\w]+)/([-_.:\w]+)/([-_.:\w]+)/runs/([-_.:\w]+)/(.*) {
    auth_request     /auth/v1/;
    auth_request_set $auth_status $upstream_status;

    resolver kube-dns.kube-system.svc.new-dns valid=5s;
    proxy_pass http://plx-operation-$4.$1.svc.new-dns;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_hide_header X-Frame-Options;
    proxy_set_header Origin "";
    proxy_set_header Host $http_host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_buffering off;
}
"""  # noqa
        settings.PROXIES_CONFIG.auth_enabled = True
        settings.PROXIES_CONFIG.dns_prefix = "kube-dns.kube-system"
        settings.PROXIES_CONFIG.dns_use_resolver = True
        settings.PROXIES_CONFIG.dns_custom_cluster = "new-dns"
        resolver = get_resolver()
        assert (
            get_services_location_config(
                resolver=resolver, auth=get_auth_config(), rewrite=False
            )
            == expected
        )

    def test_services_dns_backend(self):
        settings.PROXIES_CONFIG.auth_enabled = False
        settings.PROXIES_CONFIG.dns_use_resolver = True
        expected = r"""
location ~ /services/v1/([-_.:\w]+)/([-_.:\w]+)/([-_.:\w]+)/runs/([-_.:\w]+)/(.*) {
    resolver kube-dns.kube-system.svc.cluster.local valid=5s;
    proxy_pass http://plx-operation-$4.$1.svc.cluster.local;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_hide_header X-Frame-Options;
    proxy_set_header Origin "";
    proxy_set_header Host $http_host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_buffering off;
}
"""  # noqa
        settings.PROXIES_CONFIG.dns_custom_cluster = "cluster.local"
        assert get_dns_config() == "kube-dns.kube-system.svc.cluster.local"
        resolver = get_resolver()
        assert (
            get_services_location_config(resolver=resolver, auth="", rewrite=False)
            == expected
        )

        expected = r"""
location ~ /services/v1/([-_.:\w]+)/([-_.:\w]+)/([-_.:\w]+)/runs/([-_.:\w]+)/(.*) {
    auth_request     /auth/v1/;
    auth_request_set $auth_status $upstream_status;

    resolver kube-dns.kube-system.svc.new-dns valid=5s;
    proxy_pass http://plx-operation-$4.$1.svc.new-dns;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_hide_header X-Frame-Options;
    proxy_set_header Origin "";
    proxy_set_header Host $http_host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_buffering off;
}
"""  # noqa
        settings.PROXIES_CONFIG.auth_enabled = True
        settings.PROXIES_CONFIG.dns_custom_cluster = "new-dns"
        assert get_dns_config() == "kube-dns.kube-system.svc.new-dns"
        resolver = get_resolver()
        assert (
            get_services_location_config(
                resolver=resolver, auth=get_auth_config(), rewrite=False
            )
            == expected
        )

    def test_services_dns_prefix(self):
        settings.PROXIES_CONFIG.auth_enabled = True
        settings.PROXIES_CONFIG.dns_use_resolver = True
        expected = r"""
location ~ /services/v1/([-_.:\w]+)/([-_.:\w]+)/([-_.:\w]+)/runs/([-_.:\w]+)/(.*) {
    auth_request     /auth/v1/;
    auth_request_set $auth_status $upstream_status;

    resolver coredns.kube-system.svc.cluster.local valid=5s;
    proxy_pass http://plx-operation-$4.$1.svc.cluster.local;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_hide_header X-Frame-Options;
    proxy_set_header Origin "";
    proxy_set_header Host $http_host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_buffering off;
}
"""  # noqa
        settings.PROXIES_CONFIG.dns_prefix = "coredns.kube-system"
        settings.PROXIES_CONFIG.dns_custom_cluster = "cluster.local"
        assert get_dns_config() == "coredns.kube-system.svc.cluster.local"
        resolver = get_resolver()
        assert (
            get_services_location_config(
                resolver=resolver, auth=get_auth_config(), rewrite=False
            )
            == expected
        )

        expected = r"""
location ~ /services/v1/([-_.:\w]+)/([-_.:\w]+)/([-_.:\w]+)/runs/([-_.:\w]+)/(.*) {
    auth_request     /auth/v1/;
    auth_request_set $auth_status $upstream_status;

    resolver kube-dns.new-system.svc.new-dns valid=5s;
    proxy_pass http://plx-operation-$4.$1.svc.new-dns;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_hide_header X-Frame-Options;
    proxy_set_header Origin "";
    proxy_set_header Host $http_host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_buffering off;
}
"""  # noqa
        settings.PROXIES_CONFIG.dns_prefix = "kube-dns.new-system"
        settings.PROXIES_CONFIG.dns_custom_cluster = "new-dns"
        assert get_dns_config() == "kube-dns.new-system.svc.new-dns"
        resolver = get_resolver()
        assert (
            get_services_location_config(
                resolver=resolver, auth=get_auth_config(), rewrite=False
            )
            == expected
        )


@pytest.mark.proxies_mark
class TestGatewayRewriteServicesSchemas(BaseTestCase):
    SET_PROXIES_SETTINGS = True

    def test_service_dns_resolver(self):
        settings.PROXIES_CONFIG.auth_enabled = False
        expected = r"""
location ~ /rewrite-services/v1/([-_.:\w]+)/([-_.:\w]+)/([-_.:\w]+)/runs/([-_.:\w]+)/(.*) {
    rewrite_log on;
    rewrite ^/rewrite-services/v1/([-_.:\w]+)/([-_.:\w]+)/([-_.:\w]+)/runs/([-_.:\w]+)/(.*) /$5 break;
    proxy_pass http://plx-operation-$4.$1.svc.cluster.local;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_hide_header X-Frame-Options;
    proxy_set_header Origin "";
    proxy_set_header Host $http_host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_buffering off;
}
"""  # noqa
        settings.PROXIES_CONFIG.dns_use_resolver = False
        resolver = get_resolver()
        assert (
            get_services_location_config(resolver=resolver, auth="", rewrite=True)
            == expected
        )

        expected = r"""
location ~ /rewrite-services/v1/([-_.:\w]+)/([-_.:\w]+)/([-_.:\w]+)/runs/([-_.:\w]+)/(.*) {
    auth_request     /auth/v1/;
    auth_request_set $auth_status $upstream_status;

    resolver kube-dns.kube-system.svc.new-dns valid=5s;
    rewrite_log on;
    rewrite ^/rewrite-services/v1/([-_.:\w]+)/([-_.:\w]+)/([-_.:\w]+)/runs/([-_.:\w]+)/(.*) /$5 break;
    proxy_pass http://plx-operation-$4.$1.svc.new-dns;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_hide_header X-Frame-Options;
    proxy_set_header Origin "";
    proxy_set_header Host $http_host;
    proxy_set_header X-Real-IP $remote_addr;