|
1 | 1 | # -*- coding: utf-8 -*- |
2 | 2 | # pylint: disable=ungrouped-imports |
3 | 3 |
|
| 4 | +import enum |
4 | 5 | import os |
5 | 6 | from typing import TYPE_CHECKING |
6 | 7 |
|
|
10 | 11 |
|
11 | 12 | if TYPE_CHECKING: |
12 | 13 | from datetime import datetime |
13 | | - from typing import Any, Dict, List |
| 14 | + from enum import IntEnum |
| 15 | + from typing import Any, Dict, List, Optional |
14 | 16 |
|
15 | 17 | # database client |
16 | 18 | DB = playhouse.db_url.connect(os.getenv('DB_URL', 'mysql://127.0.0.1')) |
17 | 19 |
|
18 | 20 |
|
| 21 | +class IntEnumField(peewee.IntegerField): |
| 22 | + """:class:`enum.IntEnum` data field.""" |
| 23 | + |
| 24 | + #: The original :class:`enum.IntEnum` class. |
| 25 | + choices: 'IntEnum' |
| 26 | + |
| 27 | + # def db_value(self, value: 'Optional[IntEnum]') -> 'Optional[str]': # pylint: disable=inconsistent-return-statements |
| 28 | + # """Dump the value for database storage. |
| 29 | + |
| 30 | + # Args: |
| 31 | + # val: Original enumeration object. |
| 32 | + |
| 33 | + # Returns: |
| 34 | + # Integral representation of the enumeration. |
| 35 | + |
| 36 | + # """ |
| 37 | + # if value is not None: |
| 38 | + # return value |
| 39 | + |
| 40 | + def python_value(self, value: 'Optional[int]') -> 'Optional[IntEnum]': # pylint: disable=inconsistent-return-statements |
| 41 | + """Load the value from database storage. |
| 42 | +
|
| 43 | + Args: |
| 44 | + value: Integral representation of the enumeration. |
| 45 | +
|
| 46 | + Returns: |
| 47 | + Original enumeration object. |
| 48 | +
|
| 49 | + """ |
| 50 | + if value is not None: |
| 51 | + return self.choices(value) # type: ignore |
| 52 | + return None |
| 53 | + |
| 54 | + |
| 55 | +class Proxy(enum.IntEnum): |
| 56 | + """Proxy types supported by :mod:`darc`. |
| 57 | +
|
| 58 | + .. _tor2web: https://onion.sh/ |
| 59 | + """ |
| 60 | + |
| 61 | + #: No proxy. |
| 62 | + NULL = enum.auto() |
| 63 | + #: Tor proxy. |
| 64 | + TOR = enum.auto() |
| 65 | + #: I2P proxy. |
| 66 | + I2P = enum.auto() |
| 67 | + #: ZeroNet proxy. |
| 68 | + ZERONET = enum.auto() |
| 69 | + #: Freenet proxy. |
| 70 | + FREENET = enum.auto() |
| 71 | + #: Proxied Tor (`tor2web`_, no proxy). |
| 72 | + TOR2WEB = enum.auto() |
| 73 | + |
| 74 | + |
19 | 75 | def table_function(model_class: peewee.Model) -> str: |
20 | 76 | """Generate table name dynamically. |
21 | 77 |
|
@@ -81,9 +137,9 @@ class HostnameModel(BaseModel): |
81 | 137 | """Data model for a hostname record.""" |
82 | 138 |
|
83 | 139 | #: Hostname (c.f. :attr:`link.host <darc.link.Link.host>`). |
84 | | - hostname: str = peewee.TextField() |
| 140 | + hostname: str = peewee.CharField(max_length=255, unique=True) # a valid FQDN is at most 255 characters |
85 | 141 | #: Proxy type (c.f. :attr:`link.proxy <darc.link.Link.proxy>`). |
86 | | - proxy: str = peewee.CharField(max_length=8) |
| 142 | + proxy: 'Proxy' = IntEnumField(choices=Proxy) |
87 | 143 |
|
88 | 144 | #: Timestamp of first ``new_host`` submission. |
89 | 145 | discovery: 'datetime' = peewee.DateTimeField() |
|
0 commit comments