eshaan7 / Flask-Shell2HTTP

Execute shell commands via HTTP server (via flask's endpoints).
https://flask-shell2http.readthedocs.io/
BSD 3-Clause "New" or "Revised" License
170 stars 28 forks source link

Add some kind of basic auth #35

Closed madhavajay closed 2 years ago

madhavajay commented 2 years ago

I added some basic key checking but I wanted to know if there was a better way to add it, or if you want me to open a PR which adds some kind of hook to allow a custom function to be run on either the get or post to validate the args and in doing so allow for some kind of key / auth check?

# base_entrypoint.py
# change import to local version
from .api import Shell2HttpAPI

# api.py
# auth function
def check_stack_api_key(challenge_key: str) -> bool:
    key = os.environ.get("STACK_API_KEY", None)  # Get key from environment
    if key is None:
        return False
    if challenge_key == key:
        return True
    return False

class Shell2HttpAPI(MethodView):
    def get(self):
        ...

        # call the auth check
        stack_api_key = request.args.get("STACK_API_KEY", None)
        if not check_stack_api_key(challenge_key=stack_api_key):
            raise Exception("STACK_API_KEY doesn't match.")

    def post(self):
        ...
        # call the auth check
        json_input = request.get_json()
        if not check_stack_api_key(
            challenge_key=json_input.get("STACK_API_KEY", None)
        ):
            raise Exception("STACK_API_KEY doesn't match.")
eshaan7 commented 2 years ago

Hey!

The recommended way to add a "custom hook" is to use decorators. See following example:

https://github.com/eshaan7/Flask-Shell2HTTP/blob/1c568c9f523058205850f1fd0cf2d7899a22a47d/examples/with_decorators.py#L31-L38 (Here inside the decorator function you can access the Flask's global objects like Flask.request or Flask.g making it possible to do anything)

My recommendation:

import functools
from flask import request

 def basic_auth_check(f): 
     @functools.wraps(f) 
     def inner_decorator(*args, **kwargs): 
         stack_api_key = request.headers.get("X-STACK-API-KEY", None)
         if not check_stack_api_key(challenge_key=stack_api_key):
             raise Exception("STACK_API_KEY doesn't match.")
         return f(*args, **kwargs) 

     return inner_decorator

So with the above implementation when making the request, you should include a request header X-STACK-API-KEY.

I recommend using this custom request header method because it will remain consistent for both GET and POST calls and won't interfere with the internal working of flask_shell2http.

madhavajay commented 2 years ago

@eshaan7 Awesome example, can you add this to the examples folder for other users who might want some kind of auth?