Coverage for chatgpt_proxy / gen_api_key.py: 89%

1# MIT License 

2# 

3# Copyright (c) 2025 Tuomo Kriikkula 

4# 

5# Permission is hereby granted, free of charge, to any person obtaining a copy 

6# of this software and associated documentation files (the "Software"), to deal 

7# in the Software without restriction, including without limitation the rights 

8# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

9# copies of the Software, and to permit persons to whom the Software is 

10# furnished to do so, subject to the following conditions: 

11# 

12# The above copyright notice and this permission notice shall be included in all 

13# copies or substantial portions of the Software. 

14# 

15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

21# SOFTWARE. 

22 

23# Implements a simple proxy server for communication between an UnrealScript 

24# client and OpenAI servers. 

25 

26import asyncio 

27import datetime 

28import hashlib 

29import ipaddress 

30import os 

31 

32import asyncpg 

33import click 

34import jwt 

35from asyncpg import Connection 

36 

37from chatgpt_proxy.db import queries 

38from chatgpt_proxy.utils import utcnow 

39 

40 

41async def async_main( 

42 game_server_address: ipaddress.IPv4Address, 

43 game_server_port: int, 

44 secret: str, 

45 issuer: str, 

46 audience: str, 

47 expires_at: datetime.datetime, 

48 name: str | None = None, 

49) -> str: 

50 conn: Connection | None = None 

51 url = os.environ["DATABASE_URL"] 

52 try: 

53 iat = utcnow() 

54 token = jwt.encode( 

55 key=secret, 

56 algorithm="HS256", 

57 payload={ 

58 "iss": issuer, 

59 "aud": audience, 

60 "sub": f"{game_server_address}:{game_server_port}", 

61 "exp": int(expires_at.timestamp()), 

62 "iat": int(iat.timestamp()), 

63 }, 

64 ) 

65 token_sha256 = hashlib.sha256(token.encode("utf-8")).digest() 

66 conn = await asyncpg.connect(url) 

67 await queries.insert_game_server_api_key( 

68 conn=conn, 

69 issued_at=iat, 

70 expires_at=expires_at, 

71 token_hash=token_sha256, 

72 game_server_address=game_server_address, 

73 game_server_port=game_server_port, 

74 name=name, 

75 ) 

76 return token 

77 finally: 

78 if conn: 

79 await conn.close() 

80 

81 

82@click.command() 

83@click.option("--game-server-address", "-a", type=ipaddress.IPv4Address, required=True) 

84@click.option("--game-server-port", "-p", type=int, required=True) 

85@click.option("--issuer", "-i", type=str, required=True) 

86@click.option("--audience", "-u", type=str, required=True) 

87@click.option("--expires-at", "-e", type=float, required=True) 

88@click.option("--name", "-n", type=str, default=None) 

89def main( 

90 game_server_address: ipaddress.IPv4Address, 

91 game_server_port: int, 

92 issuer: str, 

93 audience: str, 

94 expires_at: float, 

95 name: str | None, 

96) -> None: 

97 secret = os.environ["SANIC_SECRET"] 

98 token = asyncio.run(async_main( 

99 game_server_address=game_server_address, 

100 game_server_port=game_server_port, 

101 secret=secret, 

102 issuer=issuer, 

103 audience=audience, 

104 expires_at=datetime.datetime.fromtimestamp(expires_at), 

105 name=name, 

106 )) 

107 print(token) 

108 

109 

110if __name__ == "__main__": 

111 main()