cloudflare/cloudflared
Publicmirrored from https://github.com/cloudflare/cloudflaredAvailable
component-tests/config.py
107lines · 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 | if 'tunnel' not in config: |
| 45 | config['tunnel'] = self.tunnel |
| 46 | if 'credentials-file' not in config: |
| 47 | config['credentials-file'] = self.credentials_file |
| 48 | # In some cases we want to override default ingress, such as in config tests |
| 49 | if 'ingress' not in config: |
| 50 | config['ingress'] = self.ingress |
| 51 | return config |
| 52 | |
| 53 | |
| 54 | @dataclass(frozen=True) |
| 55 | class NamedTunnelConfig(NamedTunnelBaseConfig): |
| 56 | full_config: dict = None |
| 57 | additional_config: InitVar[dict] = {} |
| 58 | |
| 59 | def __post_init__(self, additional_config): |
| 60 | # Cannot call set self.full_config because the class is frozen, instead, we can use __setattr__ |
| 61 | # https://docs.python.org/3/library/dataclasses.html#frozen-instances |
| 62 | object.__setattr__(self, 'full_config', |
| 63 | self.merge_config(additional_config)) |
| 64 | |
| 65 | def get_url(self): |
| 66 | return "https://" + self.ingress[0]['hostname'] |
| 67 | |
| 68 | def base_config(self): |
| 69 | config = self.full_config.copy() |
| 70 | |
| 71 | # removes the tunnel reference |
| 72 | del(config["tunnel"]) |
| 73 | del(config["credentials-file"]) |
| 74 | |
| 75 | return config |
| 76 | |
| 77 | def get_tunnel_id(self): |
| 78 | return self.full_config["tunnel"] |
| 79 | |
| 80 | def get_token(self): |
| 81 | creds = self.get_credentials_json() |
| 82 | token_dict = {"a": creds["AccountTag"], "t": creds["TunnelID"], "s": creds["TunnelSecret"]} |
| 83 | token_json_str = json.dumps(token_dict) |
| 84 | return base64.b64encode(token_json_str.encode('utf-8')) |
| 85 | |
| 86 | def get_credentials_json(self): |
| 87 | with open(self.credentials_file) as json_file: |
| 88 | return json.load(json_file) |
| 89 | |
| 90 | @dataclass(frozen=True) |
| 91 | class QuickTunnelConfig(BaseConfig): |
| 92 | full_config: dict = None |
| 93 | additional_config: InitVar[dict] = {} |
| 94 | |
| 95 | def __post_init__(self, additional_config): |
| 96 | # Cannot call set self.full_config because the class is frozen, instead, we can use __setattr__ |
| 97 | # https://docs.python.org/3/library/dataclasses.html#frozen-instances |
| 98 | object.__setattr__(self, 'full_config', |
| 99 | self.merge_config(additional_config)) |
| 100 | |
| 101 | @dataclass(frozen=True) |
| 102 | class ProxyDnsConfig(BaseConfig): |
| 103 | full_config = { |
| 104 | "port": PROXY_DNS_PORT, |
| 105 | "no-autoupdate": True, |
| 106 | } |
| 107 | |
| 108 | |