Skip to content

Commit cc110c5

Browse files
michael.yakmichaelyaakoby
michael.yak
authored and
michaelyaakoby
committed
Pycturator should allow disabling specifc endpoints
In order to disable pyctuator endpoints such as `/pyctuator/httptraces`, pyctuator now can be initialized with flags selecting endpoints that shouldn't be registered in SBA and should fail if queried (GET) directly. Fixes #92
1 parent bb04e43 commit cc110c5

16 files changed

+471
-279
lines changed

‎pyctuator/endpoints.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
from enum import Flag, auto
2+
3+
4+
class Endpoints(Flag):
5+
NONE = 0
6+
ENV = auto()
7+
INFO = auto()
8+
HEALTH = auto()
9+
METRICS = auto()
10+
LOGGERS = auto()
11+
THREAD_DUMP = auto()
12+
LOGFILE = auto()
13+
HTTP_TRACE = auto()

‎pyctuator/impl/aiohttp_pyctuator.py

Lines changed: 46 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from aiohttp import web
1010
from multidict import CIMultiDictProxy
1111

12+
from pyctuator.endpoints import Endpoints
1213
from pyctuator.httptrace import TraceRecord, TraceRequest, TraceResponse
1314
from pyctuator.impl import SBA_V2_CONTENT_TYPE
1415
from pyctuator.impl.pyctuator_impl import PyctuatorImpl
@@ -17,7 +18,7 @@
1718

1819
# pylint: disable=too-many-locals,unused-argument
1920
class AioHttpPyctuator(PyctuatorRouter):
20-
def __init__(self, app: web.Application, pyctuator_impl: PyctuatorImpl) -> None:
21+
def __init__(self, app: web.Application, pyctuator_impl: PyctuatorImpl, disabled_endpoints: Endpoints) -> None:
2122
super().__init__(app, pyctuator_impl)
2223

2324
custom_dumps = partial(
@@ -106,34 +107,50 @@ async def intercept_requests_and_responses(request: web.Request, handler: Callab
106107
self.pyctuator_impl.http_tracer.add_record(record=new_record)
107108
return response
108109

109-
app.add_routes(
110-
[
111-
web.get("/pyctuator", get_endpoints),
112-
web.options("/pyctuator/env", empty_handler),
113-
web.options("/pyctuator/info", empty_handler),
114-
web.options("/pyctuator/health", empty_handler),
115-
web.options("/pyctuator/metrics", empty_handler),
116-
web.options("/pyctuator/loggers", empty_handler),
117-
web.options("/pyctuator/dump", empty_handler),
118-
web.options("/pyctuator/threaddump", empty_handler),
119-
web.options("/pyctuator/logfile", empty_handler),
120-
web.options("/pyctuator/trace", empty_handler),
121-
web.options("/pyctuator/httptrace", empty_handler),
122-
web.get("/pyctuator/env", get_environment),
123-
web.get("/pyctuator/info", get_info),
124-
web.get("/pyctuator/health", get_health),
125-
web.get("/pyctuator/metrics", get_metric_names),
126-
web.get("/pyctuator/metrics/{metric_name}", get_metric_measurement),
127-
web.get("/pyctuator/loggers", get_loggers),
128-
web.get("/pyctuator/loggers/{logger_name}", get_logger),
129-
web.post("/pyctuator/loggers/{logger_name}", set_logger_level),
130-
web.get("/pyctuator/dump", get_thread_dump),
131-
web.get("/pyctuator/threaddump", get_thread_dump),
132-
web.get("/pyctuator/logfile", get_logfile),
133-
web.get("/pyctuator/trace", get_httptrace),
134-
web.get("/pyctuator/httptrace", get_httptrace),
135-
]
136-
)
110+
routes = [
111+
web.get("/pyctuator", get_endpoints),
112+
]
113+
114+
if Endpoints.ENV not in disabled_endpoints:
115+
routes.append(web.options("/pyctuator/env", empty_handler))
116+
routes.append(web.get("/pyctuator/env", get_environment))
117+
118+
if Endpoints.INFO not in disabled_endpoints:
119+
routes.append(web.options("/pyctuator/info", empty_handler))
120+
routes.append(web.get("/pyctuator/info", get_info))
121+
122+
if Endpoints.HEALTH not in disabled_endpoints:
123+
routes.append(web.options("/pyctuator/health", empty_handler))
124+
routes.append(web.get("/pyctuator/health", get_health))
125+
126+
if Endpoints.METRICS not in disabled_endpoints:
127+
routes.append(web.options("/pyctuator/metrics", empty_handler))
128+
routes.append(web.get("/pyctuator/metrics", get_metric_names))
129+
routes.append(web.get("/pyctuator/metrics/{metric_name}", get_metric_measurement))
130+
131+
if Endpoints.LOGGERS not in disabled_endpoints:
132+
routes.append(web.options("/pyctuator/loggers", empty_handler))
133+
routes.append(web.get("/pyctuator/loggers", get_loggers))
134+
routes.append(web.get("/pyctuator/loggers/{logger_name}", get_logger))
135+
routes.append(web.post("/pyctuator/loggers/{logger_name}", set_logger_level))
136+
137+
if Endpoints.THREAD_DUMP not in disabled_endpoints:
138+
routes.append(web.options("/pyctuator/dump", empty_handler))
139+
routes.append(web.options("/pyctuator/threaddump", empty_handler))
140+
routes.append(web.get("/pyctuator/dump", get_thread_dump))
141+
routes.append(web.get("/pyctuator/threaddump", get_thread_dump))
142+
143+
if Endpoints.LOGFILE not in disabled_endpoints:
144+
routes.append(web.options("/pyctuator/logfile", empty_handler))
145+
routes.append(web.get("/pyctuator/logfile", get_logfile))
146+
147+
if Endpoints.HTTP_TRACE not in disabled_endpoints:
148+
routes.append(web.options("/pyctuator/trace", empty_handler))
149+
routes.append(web.options("/pyctuator/httptrace", empty_handler))
150+
routes.append(web.get("/pyctuator/trace", get_httptrace))
151+
routes.append(web.get("/pyctuator/httptrace", get_httptrace))
152+
153+
app.add_routes(routes)
137154
app.middlewares.append(intercept_requests_and_responses)
138155

139156
def _custom_json_serializer(self, value: Any) -> Any:

‎pyctuator/impl/fastapi_pyctuator.py

Lines changed: 70 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from starlette.requests import Request
1111
from starlette.responses import Response
1212

13+
from pyctuator.endpoints import Endpoints
1314
from pyctuator.environment.environment_provider import EnvironmentData
1415
from pyctuator.httptrace import TraceRecord, TraceRequest, TraceResponse
1516
from pyctuator.httptrace.http_tracer import Traces
@@ -33,8 +34,9 @@ def __init__(
3334
self,
3435
app: FastAPI,
3536
pyctuator_impl: PyctuatorImpl,
36-
include_in_openapi_schema: bool = False,
37-
customizer: Optional[Callable[[APIRouter], None]] = None
37+
include_in_openapi_schema: bool,
38+
customizer: Optional[Callable[[APIRouter], None]],
39+
disabled_endpoints: Endpoints,
3840
) -> None:
3941
super().__init__(app, pyctuator_impl)
4042
router = APIRouter()
@@ -64,70 +66,78 @@ def options() -> None:
6466
documentation.
6567
"""
6668

67-
@router.get("/env", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
68-
def get_environment() -> EnvironmentData:
69-
return pyctuator_impl.get_environment()
69+
if Endpoints.ENV not in disabled_endpoints:
70+
@router.get("/env", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
71+
def get_environment() -> EnvironmentData:
72+
return pyctuator_impl.get_environment()
7073

71-
@router.get("/info", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
72-
def get_info() -> Dict:
73-
return pyctuator_impl.get_app_info()
74+
if Endpoints.INFO not in disabled_endpoints:
75+
@router.get("/info", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
76+
def get_info() -> Dict:
77+
return pyctuator_impl.get_app_info()
7478

75-
@router.get("/health", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
76-
def get_health(response: Response) -> object:
77-
health = pyctuator_impl.get_health()
78-
response.status_code = health.http_status()
79-
return health
79+
if Endpoints.HEALTH not in disabled_endpoints:
80+
@router.get("/health", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
81+
def get_health(response: Response) -> object:
82+
health = pyctuator_impl.get_health()
83+
response.status_code = health.http_status()
84+
return health
8085

81-
@router.get("/metrics", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
82-
def get_metric_names() -> MetricNames:
83-
return pyctuator_impl.get_metric_names()
86+
if Endpoints.METRICS not in disabled_endpoints:
87+
@router.get("/metrics", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
88+
def get_metric_names() -> MetricNames:
89+
return pyctuator_impl.get_metric_names()
8490

85-
@router.get("/metrics/{metric_name}", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
86-
def get_metric_measurement(metric_name: str) -> Metric:
87-
return pyctuator_impl.get_metric_measurement(metric_name)
91+
@router.get("/metrics/{metric_name}", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
92+
def get_metric_measurement(metric_name: str) -> Metric:
93+
return pyctuator_impl.get_metric_measurement(metric_name)
8894

8995
# Retrieving All Loggers
90-
@router.get("/loggers", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
91-
def get_loggers() -> LoggersData:
92-
return pyctuator_impl.logging.get_loggers()
93-
94-
@router.post("/loggers/{logger_name}", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
95-
def set_logger_level(item: FastApiLoggerItem, logger_name: str) -> Dict:
96-
pyctuator_impl.logging.set_logger_level(logger_name, item.configuredLevel)
97-
return {}
98-
99-
@router.get("/loggers/{logger_name}", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
100-
def get_logger(logger_name: str) -> LoggerLevels:
101-
return pyctuator_impl.logging.get_logger(logger_name)
102-
103-
@router.get("/dump", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
104-
@router.get("/threaddump", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
105-
def get_thread_dump() -> ThreadDump:
106-
return pyctuator_impl.get_thread_dump()
107-
108-
@router.get("/logfile", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
109-
def get_logfile(range_header: str = Header(default=None,
110-
alias="range")) -> Response: # pylint: disable=redefined-builtin
111-
if not range_header:
112-
return Response(content=pyctuator_impl.logfile.log_messages.get_range())
113-
114-
str_res, start, end = pyctuator_impl.logfile.get_logfile(range_header)
115-
116-
my_res = Response(
117-
status_code=HTTPStatus.PARTIAL_CONTENT.value,
118-
content=str_res,
119-
headers={
120-
"Content-Type": "text/html; charset=UTF-8",
121-
"Accept-Ranges": "bytes",
122-
"Content-Range": f"bytes {start}-{end}/{end}",
123-
})
124-
125-
return my_res
126-
127-
@router.get("/trace", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
128-
@router.get("/httptrace", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
129-
def get_httptrace() -> Traces:
130-
return pyctuator_impl.http_tracer.get_httptrace()
96+
if Endpoints.LOGGERS not in disabled_endpoints:
97+
@router.get("/loggers", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
98+
def get_loggers() -> LoggersData:
99+
return pyctuator_impl.logging.get_loggers()
100+
101+
@router.post("/loggers/{logger_name}", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
102+
def set_logger_level(item: FastApiLoggerItem, logger_name: str) -> Dict:
103+
pyctuator_impl.logging.set_logger_level(logger_name, item.configuredLevel)
104+
return {}
105+
106+
@router.get("/loggers/{logger_name}", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
107+
def get_logger(logger_name: str) -> LoggerLevels:
108+
return pyctuator_impl.logging.get_logger(logger_name)
109+
110+
if Endpoints.THREAD_DUMP not in disabled_endpoints:
111+
@router.get("/dump", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
112+
@router.get("/threaddump", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
113+
def get_thread_dump() -> ThreadDump:
114+
return pyctuator_impl.get_thread_dump()
115+
116+
if Endpoints.LOGFILE not in disabled_endpoints:
117+
@router.get("/logfile", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
118+
def get_logfile(range_header: str = Header(default=None,
119+
alias="range")) -> Response: # pylint: disable=redefined-builtin
120+
if not range_header:
121+
return Response(content=pyctuator_impl.logfile.log_messages.get_range())
122+
123+
str_res, start, end = pyctuator_impl.logfile.get_logfile(range_header)
124+
125+
my_res = Response(
126+
status_code=HTTPStatus.PARTIAL_CONTENT.value,
127+
content=str_res,
128+
headers={
129+
"Content-Type": "text/html; charset=UTF-8",
130+
"Accept-Ranges": "bytes",
131+
"Content-Range": f"bytes {start}-{end}/{end}",
132+
})
133+
134+
return my_res
135+
136+
if Endpoints.HTTP_TRACE not in disabled_endpoints:
137+
@router.get("/trace", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
138+
@router.get("/httptrace", include_in_schema=include_in_openapi_schema, tags=["pyctuator"])
139+
def get_httptrace() -> Traces:
140+
return pyctuator_impl.http_tracer.get_httptrace()
131141

132142
@app.middleware("http")
133143
async def intercept_requests_and_responses(

0 commit comments

Comments
 (0)