cloudflare/cloudflared
Publicmirrored from https://github.com/cloudflare/cloudflaredAvailable
component-tests/config.py
128lines · modecode
| 1 | #!/usr/bin/env python |
| 2 | import copy |
| 3 | import json |
| 4 | import base64 |
| 5 | |
| 6 | from dataclasses import dataclass, InitVar |
| 7 | |
| 8 | from constants import METRICS_PORT, PROXY_DNS_PORT |
| 9 | |
| 10 | # frozen=True raises exception when assigning to fields. This emulates immutability |
| 11 | |
| 12 | |
| 13 | @dataclass(frozen=True) |
| 14 | class BaseConfig: |
| 15 | cloudflared_binary: str |
| 16 | no_autoupdate: bool = True |
| 17 | metrics: str = f'localhost:{METRICS_PORT}' |
| 18 | |
| 19 | def merge_config(self, additional): |
| 20 | config = copy.copy(additional) |
| 21 | config['no-autoupdate'] = self.no_autoupdate |
| 22 | config['metrics'] = self.metrics |
| 23 | return config |
| 24 | |
| 25 | |
| 26 | @dataclass(frozen=True) |
| 27 | class NamedTunnelBaseConfig(BaseConfig): |
| 28 | # The attributes of the parent class are ordered before attributes in this class, |
| 29 | # so we have to use default values here and check if they are set in __post_init__ |
| 30 | tunnel: str = None |
| 31 | credentials_file: str = None |
| 32 | ingress: list = None |
| 33 | |
| 34 | def __post_init__(self): |
| 35 | if self.tunnel is None: |
| 36 | raise TypeError("Field tunnel is not set") |
| 37 | if self.credentials_file is None: |
| 38 | raise TypeError("Field credentials_file is not set") |
| 39 | if self.ingress is None: |
| 40 | raise TypeError("Field ingress is not set") |
| 41 | |
| 42 | def merge_config(self, additional): |
| 43 | config = super(NamedTunnelBaseConfig, self).merge_config(additional) |
| 44 | config['tunnel'] = self.tunnel |
| 45 | config['credentials-file'] = self.credentials_file |
| 46 | # In some cases we want to override default ingress, such as in config tests |
| 47 | if 'ingress' not in config: |
| 48 | config['ingress'] = self.ingress |
| 49 | return config |
| 50 | |
| 51 | |
| 52 | @dataclass(frozen=True) |
| 53 | class NamedTunnelConfig(NamedTunnelBaseConfig): |
| 54 | full_config: dict = None |
| 55 | additional_config: InitVar[dict] = {} |
| 56 | |
| 57 | def __post_init__(self, additional_config): |
| 58 | # Cannot call set self.full_config because the class is frozen, instead, we can use __setattr__ |
| 59 | # https://docs.python.org/3/library/dataclasses.html#frozen-instances |
| 60 | object.__setattr__(self, 'full_config', |
| 61 | self.merge_config(additional_config)) |
| 62 | |
| 63 | def get_url(self): |
| 64 | return "https://" + self.ingress[0]['hostname'] |
| 65 | |
| 66 | def base_config(self): |
| 67 | config = self.full_config.copy() |
| 68 | |
| 69 | # removes the tunnel reference |
| 70 | del(config["tunnel"]) |
| 71 | del(config["credentials-file"]) |
| 72 | |
| 73 | return config |
| 74 | |
| 75 | def get_tunnel_id(self): |
| 76 | return self.full_config["tunnel"] |
| 77 | |
| 78 | def get_token(self): |
| 79 | creds = self.get_credentials_json() |
| 80 | token_dict = {"a": creds["AccountTag"], "t": creds["TunnelID"], "s": creds["TunnelSecret"]} |
| 81 | token_json_str = json.dumps(token_dict) |
| 82 | return base64.b64encode(token_json_str.encode('utf-8')) |
| 83 | |
| 84 | def get_credentials_json(self): |
| 85 | with open(self.credentials_file) as json_file: |
| 86 | return json.load(json_file) |
| 87 | |
| 88 | |
| 89 | @dataclass(frozen=True) |
| 90 | class ClassicTunnelBaseConfig(BaseConfig): |
| 91 | hostname: str = None |
| 92 | origincert: str = None |
| 93 | |
| 94 | def __post_init__(self): |
| 95 | if self.hostname is None: |
| 96 | raise TypeError("Field tunnel is not set") |
| 97 | if self.origincert is None: |
| 98 | raise TypeError("Field credentials_file is not set") |
| 99 | |
| 100 | def merge_config(self, additional): |
| 101 | config = super(ClassicTunnelBaseConfig, self).merge_config(additional) |
| 102 | config['hostname'] = self.hostname |
| 103 | config['origincert'] = self.origincert |
| 104 | return config |
| 105 | |
| 106 | |
| 107 | @dataclass(frozen=True) |
| 108 | class ClassicTunnelConfig(ClassicTunnelBaseConfig): |
| 109 | full_config: dict = None |
| 110 | additional_config: InitVar[dict] = {} |
| 111 | |
| 112 | def __post_init__(self, additional_config): |
| 113 | # Cannot call set self.full_config because the class is frozen, instead, we can use __setattr__ |
| 114 | # https://docs.python.org/3/library/dataclasses.html#frozen-instances |
| 115 | object.__setattr__(self, 'full_config', |
| 116 | self.merge_config(additional_config)) |
| 117 | |
| 118 | def get_url(self): |
| 119 | return "https://" + self.hostname |
| 120 | |
| 121 | |
| 122 | @dataclass(frozen=True) |
| 123 | class ProxyDnsConfig(BaseConfig): |
| 124 | full_config = { |
| 125 | "port": PROXY_DNS_PORT, |
| 126 | "no-autoupdate": True, |
| 127 | } |
| 128 | |
| 129 | |