commit f38e45c2e0ea15c1882308299fbe24f6c46b8243 · pyrox.dev/nixpkgs

+57 -20

nixos/tests/google-oslogin/server.py

···

       7
        
       import base64

     

       8
        
       

     

       9
        
       from http.server import BaseHTTPRequestHandler, HTTPServer

     

       0
        
       
     

       10
        
       from typing import Dict

     

       11
        
       

     

       12
        
       SNAKEOIL_PUBLIC_KEY = os.environ['SNAKEOIL_PUBLIC_KEY']

     

       13
        
       

     

       14
        
       

     

       15
       -
       def w(msg):

     

       16
        
           sys.stderr.write(f"{msg}\n")

     

       17
        
           sys.stderr.flush()

     

       18
        
       

     

       19
        
       

     

       20
       -
       def gen_fingerprint(pubkey):

     

       21
        
           decoded_key = base64.b64decode(pubkey.encode("ascii").split()[1])

     

       22
        
           return hashlib.sha256(decoded_key).hexdigest()

     

       23
        
       

     

       24
       -
       def gen_email(username):

     

       0
        
       
     

       25
        
           """username seems to be a 21 characters long number string, so mimic that in a reproducible way"""

     

       26
        
           return str(int(hashlib.sha256(username.encode()).hexdigest(), 16))[0:21]

     

       0
        
       
     

       27
        
       

     

       28
        
       def gen_mockuser(username: str, uid: str, gid: str, home_directory: str, snakeoil_pubkey: str) -> Dict:

     

       29
        
           snakeoil_pubkey_fingerprint = gen_fingerprint(snakeoil_pubkey)

     
···

       56
        
       

     

       57
        
       

     

       58
        
       class ReqHandler(BaseHTTPRequestHandler):

     

       59
       -
           def _send_json_ok(self, data):

     

       0
        
       
     

       60
        
               self.send_response(200)

     

       61
        
               self.send_header('Content-type', 'application/json')

     

       62
        
               self.end_headers()

     
···

       64
        
               w(out)

     

       65
        
               self.wfile.write(out)

     

       66
        
       

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       67
        
           def do_GET(self):

     

       68
        
               p = str(self.path)

     

       69
       -
               # mockuser and mockadmin are allowed to login, both use the same snakeoil public key

     

       70
       -
               if p == '/computeMetadata/v1/oslogin/users?username=mockuser' \

     

       71
       -
                   or p == '/computeMetadata/v1/oslogin/users?uid=1009719690':

     

       72
       -
                   self._send_json_ok(gen_mockuser(username='mockuser', uid='1009719690', gid='1009719690',

     

       73
       -
                                                   home_directory='/home/mockuser', snakeoil_pubkey=SNAKEOIL_PUBLIC_KEY))

     

       74
       -
               elif p == '/computeMetadata/v1/oslogin/users?username=mockadmin' \

     

       75
       -
                   or p == '/computeMetadata/v1/oslogin/users?uid=1009719691':

     

       76
       -
                   self._send_json_ok(gen_mockuser(username='mockadmin', uid='1009719691', gid='1009719691',

     

       77
       -
                                                   home_directory='/home/mockadmin', snakeoil_pubkey=SNAKEOIL_PUBLIC_KEY))

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       78
        
       

     

       79
       -
               # mockuser is allowed to login

     

       80
       -
               elif p == f"/computeMetadata/v1/oslogin/authorize?email={gen_email('mockuser')}&policy=login":

     

       81
       -
                   self._send_json_ok({'success': True})

     

       82
        
       

     

       83
       -
               # mockadmin may also become root

     

       84
       -
               elif p == f"/computeMetadata/v1/oslogin/authorize?email={gen_email('mockadmin')}&policy=login" or p == f"/computeMetadata/v1/oslogin/authorize?email={gen_email('mockadmin')}&policy=adminLogin":

     

       85
       -
                   self._send_json_ok({'success': True})

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       86
        
               else:

     

       87
        
                   sys.stderr.write(f"Unhandled path: {p}\n")

     

       88
        
                   sys.stderr.flush()

     

       89
       -
                   self.send_response(501)

     

       90
        
                   self.end_headers()

     

       91
        
                   self.wfile.write(b'')

     

       92

···

       7
        
       import base64

     

       8
        
       

     

       9
        
       from http.server import BaseHTTPRequestHandler, HTTPServer

     

       10
       +
       from urllib.parse import urlparse, parse_qs

     

       11
        
       from typing import Dict

     

       12
        
       

     

       13
        
       SNAKEOIL_PUBLIC_KEY = os.environ['SNAKEOIL_PUBLIC_KEY']

     

       14
        
       

     

       15
        
       

     

       16
       +
       def w(msg: bytes):

     

       17
        
           sys.stderr.write(f"{msg}\n")

     

       18
        
           sys.stderr.flush()

     

       19
        
       

     

       20
        
       

     

       21
       +
       def gen_fingerprint(pubkey: str):

     

       22
        
           decoded_key = base64.b64decode(pubkey.encode("ascii").split()[1])

     

       23
        
           return hashlib.sha256(decoded_key).hexdigest()

     

       24
        
       

     

       25
       +
       

     

       26
       +
       def gen_email(username: str):

     

       27
        
           """username seems to be a 21 characters long number string, so mimic that in a reproducible way"""

     

       28
        
           return str(int(hashlib.sha256(username.encode()).hexdigest(), 16))[0:21]

     

       29
       +
       

     

       30
        
       

     

       31
        
       def gen_mockuser(username: str, uid: str, gid: str, home_directory: str, snakeoil_pubkey: str) -> Dict:

     

       32
        
           snakeoil_pubkey_fingerprint = gen_fingerprint(snakeoil_pubkey)

     
···

       59
        
       

     

       60
        
       

     

       61
        
       class ReqHandler(BaseHTTPRequestHandler):

     

       62
       +
       

     

       63
       +
           def _send_json_ok(self, data: dict):

     

       64
        
               self.send_response(200)

     

       65
        
               self.send_header('Content-type', 'application/json')

     

       66
        
               self.end_headers()

     
···

       68
        
               w(out)

     

       69
        
               self.wfile.write(out)

     

       70
        
       

     

       71
       +
           def _send_json_success(self, success=True):

     

       72
       +
               self.send_response(200)

     

       73
       +
               self.send_header('Content-type', 'application/json')

     

       74
       +
               self.end_headers()

     

       75
       +
               out = json.dumps({"success": success}).encode()

     

       76
       +
               w(out)

     

       77
       +
               self.wfile.write(out)

     

       78
       +
       

     

       79
       +
           def _send_404(self):

     

       80
       +
               self.send_response(404)

     

       81
       +
               self.end_headers()

     

       82
       +
       

     

       83
        
           def do_GET(self):

     

       84
        
               p = str(self.path)

     

       85
       +
               pu = urlparse(p)

     

       86
       +
               params = parse_qs(pu.query)

     

       87
       +
       

     

       88
       +
               # users endpoint

     

       89
       +
               if pu.path == "/computeMetadata/v1/oslogin/users":

     

       90
       +
                   # mockuser and mockadmin are allowed to login, both use the same snakeoil public key

     

       91
       +
                   if params.get('username') == ['mockuser'] or params.get('uid') == ["1009719690"]:

     

       92
       +
                       username = "mockuser"

     

       93
       +
                       uid = "1009719690"

     

       94
       +
                   elif params.get('username') == ['mockadmin'] or params.get('uid') == ["1009719691"]:

     

       95
       +
                       username = "mockadmin"

     

       96
       +
                       uid = "1009719691"

     

       97
       +
                   else:

     

       98
       +
                       self._send_404()

     

       99
       +
                       return

     

       100
        
       

     

       101
       +
                   self._send_json_ok(gen_mockuser(username=username, uid=uid, gid=uid, home_directory=f"/home/{username}", snakeoil_pubkey=SNAKEOIL_PUBLIC_KEY))

     

       102
       +
                   return

     

       0
        
       
     

       103
        
       

     

       104
       +
               # authorize endpoint

     

       105
       +
               elif pu.path == "/computeMetadata/v1/oslogin/authorize":

     

       106
       +
                   # is user allowed to login?

     

       107
       +
                   if params.get("policy") == ["login"]:

     

       108
       +
                       # mockuser and mockadmin are allowed to login

     

       109
       +
                       if params.get('email') == [gen_email("mockuser")] or params.get('email') == [gen_email("mockadmin")]:

     

       110
       +
                           self._send_json_success()

     

       111
       +
                           return

     

       112
       +
                       self._send_json_success(False)

     

       113
       +
                       return

     

       114
       +
                   # is user allowed to become root?

     

       115
       +
                   elif params.get("policy") == ["adminLogin"]:

     

       116
       +
                       # only mockadmin is allowed to become admin

     

       117
       +
                       self._send_json_success((params['email'] == [gen_email("mockadmin")]))

     

       118
       +
                       return

     

       119
       +
                   # send 404 for other policies

     

       120
       +
                   else:

     

       121
       +
                       self._send_404()

     

       122
       +
                       return

     

       123
        
               else:

     

       124
        
                   sys.stderr.write(f"Unhandled path: {p}\n")

     

       125
        
                   sys.stderr.flush()

     

       126
       +
                   self.send_response(404)

     

       127
        
                   self.end_headers()

     

       128
        
                   self.wfile.write(b'')

     

       129