Created
May 29, 2017 07:15
-
-
Save luc-lynx/6d0fbf16c587bafeffeb15a34186d320 to your computer and use it in GitHub Desktop.
grpc custom authentication scheme example (python)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import grpc | |
import helloworld_pb2 | |
import helloworld_pb2_grpc | |
# you need to use secure port, | |
# otherwise call credentials won't be transmitted | |
def run(): | |
with open('server.crt', 'rb') as f: | |
trusted_certs = f.read() | |
credentials = grpc.ssl_channel_credentials(root_certificates=trusted_certs) | |
# make sure that all headers are in lowecase, otherwise grpc throws an exception | |
call_credentials = grpc.metadata_call_credentials( | |
lambda context, callback: callback((("x-custom-token", "SampleToken"),), None)) | |
# use this if you want standard "Authorization" header | |
#call_credentials = grpc.access_token_call_credentials("test_access_token") | |
composite_credentials = grpc.composite_channel_credentials(credentials, call_credentials) | |
channel = grpc.secure_channel('{}:{}'.format('localhost', 50051), composite_credentials) | |
stub = helloworld_pb2_grpc.GreeterStub(channel) | |
response = stub.SayHello(helloworld_pb2.HelloRequest(name='Evgeny')) | |
print("[~] Greeter client received: {}".format(response.message)) | |
response = stub.SayHelloAgain(helloworld_pb2.HelloRequest(name='Test')) | |
print("[~] Greeter client received: {}".format(response.message)) | |
if __name__ == '__main__': | |
run() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import grpc | |
import helloworld_pb2 | |
import helloworld_pb2_grpc | |
class CheckToken(object): | |
def __init__(self, empty_response_cls, grpc_status_code=grpc.StatusCode.UNAUTHENTICATED, grpc_details=""): | |
self._empty_rsp_cls = empty_response_cls | |
self._status = grpc_status_code | |
self._details = grpc_details | |
def __call__(self, f): | |
def wrapped_function(slf, request, context): | |
meta = context.invocation_metadata() | |
for item in meta: | |
if item[0] == "x-custom-token" and item[1] == "SampleToken": | |
# use this if you want to update the context and pass some data to the method handler | |
# context.user_by_token = "sample_user" | |
return f(slf, request, context) | |
context.set_code(self._status) | |
context.set_details(self._details) | |
return self._empty_rsp_cls() | |
return wrapped_function | |
class Greeter(helloworld_pb2_grpc.GreeterServicer): | |
@CheckToken(helloworld_pb2.HelloReply, grpc.StatusCode.UNAUTHENTICATED, "Invalid token") | |
def SayHello(self, request, context): | |
# context.user_by_token is available here | |
print("[~] SayHello call") | |
return helloworld_pb2.HelloReply(message='Hello, {}!'.format(request.name)) | |
def SayHelloAgain(self, request, context): | |
print("[~] SayHelloAgain call") | |
return helloworld_pb2.HelloReply(message="Hello again, {}!".format(request.name)) | |
def serve(): | |
with open('server.key', 'rb') as f: | |
private_key = f.read() | |
with open('server.crt', 'rb') as f: | |
certificate_chain = f.read() | |
server_credentials = grpc.ssl_server_credentials(((private_key, certificate_chain,),)) | |
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) | |
helloworld_pb2_grpc.add_GreeterServicer_to_server(Greeter(), server) | |
server.add_secure_port('[::]:50051', server_credentials) | |
server.start() | |
try: | |
while True: | |
time.sleep(_ONE_DAY_IN_SECONDS) | |
except KeyboardInterrupt: | |
server.stop(0) | |
if __name__ == '__main__': | |
serve() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment