Azure allowlist manager (#1222)

src/ClusterBootstrap/ctl.py

@@ -35,11 +35,25 @@
 from utils import walk_json, RestUtil
 
 
+def add_azure_params(config):
+    if "azure_cluster" not in config:
+        return
+
+    if "nsg_name" not in config["azure_cluster"]:
+        config["azure_cluster"]["nsg_name"] = config.get(
+            "nsg_name", config["cluster_name"] + "-nsg")
+
+    if "resource_group" not in config["azure_cluster"]:
+        config["azure_cluster"]["resource_group"] = config[
+            "azure_cluster"]["cluster_name"] + "ResGrp"
+
+
 def load_config_4_ctl(args, command):
     # if we need to load all config
     if command in ["svc", "render_template", "download", "docker", "db", "quota"]:
         args.config = [ENV_CNF_YAML, STATUS_YAML] if not args.config else args.config
         config = load_deploy_config(args)
+        add_azure_params(config)
     else:
         if not args.config and command != "restorefromdir":
             args.config = [STATUS_YAML]

src/ClusterBootstrap/params.py

@@ -111,6 +111,8 @@
     "storagemanager": "storagemanager",
     "repairmanager": "repairmanager",
     "repairmanageragent": "repairmanageragent",
+    # AllowList Manager mapping
+    "allowlistmanager": "allowlist-manager",
     "ssh_cert": "./deploy/sshkey/id_rsa",
     "admin_username": "core",
     # the path of where dfs/nfs is source linked and consumed on each node,
@@ -305,6 +307,7 @@
         "user-synchronizer": "etcd_node_1",
         "job-insighter": "etcd_node_1",
         "dashboard": "etcd_node_1",
+        "allowlist-manager": "etcd_node_1",
     },
     "default_kube_labels_by_node_role": {
         'infra': {

src/ClusterBootstrap/services/allowlist-manager/allowlist-manager.yaml

@@ -0,0 +1,71 @@
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: allowlist-manager-configmap
+  namespace: kube-system
+data:
+  config.yaml: |-
+    subscription: {{cnf["azure_cluster"]["subscription"]}}
+    resource_group: {{cnf["azure_cluster"]["resource_group"]}}
+    nsg_name: {{cnf["azure_cluster"]["nsg_name"]}}
+    tenant_id: {{cnf["admin-service-principal"]["tenant_id"]}}
+    client_id: {{cnf["admin-service-principal"]["client_id"]}}
+    password: {{cnf["admin-service-principal"]["password"]}}
+---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: allowlist-manager
+  namespace: kube-system
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: allowlist-manager
+  template:
+    metadata:
+      name: allowlist-manager
+      labels:
+        task: access-control
+        app: allowlist-manager
+    spec:
+      hostNetwork: true
+      nodeSelector:
+        allowlist-manager: active
+      containers:
+      - name: allowlist-manager
+        image: {{cnf["worker-dockerregistry"]}}/{{cnf["dockerprefix"]}}/{{cnf["allowlistmanager"]}}:{{cnf["dockertag"]}}
+        command:
+          - "python3"
+          - "/allowlist-manager/main.py"
+        args:
+          - '--config=/etc/allowlist-manager'
+          - '--log=/var/log/allowlist-manager'
+        imagePullPolicy: Always
+        volumeMounts:
+        - mountPath: /etc/allowlist-manager
+          name: config-volume
+        - mountPath: /var/log/allowlist-manager
+          name: log
+      {% if cnf["private_docker_registry_username"] %}
+      imagePullSecrets:
+      - name: svccred
+      {% endif %}
+      volumes:
+      - name: config-volume
+        configMap:
+          name: allowlist-manager-configmap
+      - name: log
+        hostPath:
+          path: /var/log/allowlist-manager
+      tolerations:
+      - key: CriticalAddonsOnly
+        operator: Exists
+      - key: node-role.kubernetes.io/master
+        effect: NoSchedule
+      - key: node.kubernetes.io/memory-pressure
+        operator: "Exists"
+      - key: node.kubernetes.io/disk-pressure
+        operator: "Exists"
+      - key: node-role.kubernetes.io/master
+        operator: "Exists"

src/ClusterBootstrap/services/allowlist-manager/launch_order

@@ -0,0 +1 @@
+allowlist-manager.yaml

src/ClusterBootstrap/services/allowlist-manager/pre-render.sh

@@ -0,0 +1 @@
+#!/bin/bash

src/ClusterManager/init_db.py

@@ -165,6 +165,21 @@
 );
 """)
 
+cursor.execute("""
+CREATE TABLE IF NOT EXISTS `allowlist`
+(
+    `id`           INT          NOT NULL AUTO_INCREMENT,
+    `user`         VARCHAR(255) NOT NULL UNIQUE,
+    `ip`           VARCHAR(255) NOT NULL,
+    `valid_util`   DATETIME     NOT NULL,
+    `time`         DATETIME     NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    PRIMARY KEY (`id`),
+    UNIQUE(`user`),
+    INDEX(`user`),
+    INDEX(`time`)
+);
+""")
+
 cursor.close()
 conn.commit()
 

src/ClusterManager/test/test_admin_ops.py

@@ -47,3 +47,35 @@ def test_vc_quota_change(args):
         assert gpu_memory == int(nodes * gpu_memory_per_node)
         assert cpu == int(nodes * cpu_per_node)
         assert memory == int(nodes * memory_per_node)
+
+
+@utils.case()
+def test_allow_records(args):
+    with utils.AllowRecord(args.rest, args.email, "test_user1", "10.0.0.1"):
+        with utils.AllowRecord(
+                args.rest, args.email, "test_user2", "10.0.0.2"):
+
+            # Update test_user1 record
+            resp = utils.add_allow_record(
+                args.rest, args.email, "test_user1", "10.0.0.3")
+            assert resp.status_code == 200
+
+            resp = utils.get_allow_record(args.rest, args.email, "all")
+            assert resp.status_code == 200
+            allow_records = resp.json()
+            assert "test_user1" in [record["user"] for record in allow_records]
+            assert "test_user2" in [record["user"] for record in allow_records]
+
+            for record in allow_records:
+                if record["user"] == "test_user1":
+                    assert record["ip"] == "10.0.0.3"
+                elif record["user"] == "test_user2":
+                    assert record["ip"] == "10.0.0.2"
+
+    resp = utils.get_allow_record(args.rest, args.email, "all")
+    assert resp.status_code == 200
+    allow_records = resp.json()
+    assert "test_user1" not in [record["user"] for record in allow_records]
+    assert "test_user2" not in [record["user"] for record in allow_records]
+
+

src/ClusterManager/test/utils.py

@@ -611,6 +611,37 @@ def get_active_jobs(rest_url):
     return resp.json()
 
 
+def get_allow_record(rest_url, email, user):
+    args = urllib.parse.urlencode({
+        "userName": email,
+        "user": user,
+    })
+    url = urllib.parse.urljoin(rest_url, "/AllowRecord") + "?" + args
+    resp = requests.get(url, timeout=10)
+    return resp
+
+
+def add_allow_record(rest_url, email, user, ip):
+    args = urllib.parse.urlencode({
+        "userName": email,
+        "user": user,
+        "ip": ip,
+    })
+    url = urllib.parse.urljoin(rest_url, "/AllowRecord") + "?" + args
+    resp = requests.post(url, timeout=10)
+    return resp
+
+
+def delete_allow_record(rest_url, email, user):
+    args = urllib.parse.urlencode({
+        "userName": email,
+        "user": user,
+    })
+    url = urllib.parse.urljoin(rest_url, "/AllowRecord") + "?" + args
+    resp = requests.delete(url, timeout=10)
+    return resp
+
+
 class run_job(object):
     def __init__(self, rest_url, job_spec):
         self.rest_url = rest_url
@@ -700,6 +731,27 @@ def __exit__(self, type, value, traceback):
                              json.dumps(self.origin_quota_spec))
 
 
+class AllowRecord(object):
+    def __init__(self, rest_url, username, user, ip):
+        self.rest_url = rest_url
+        self.username = username
+        self.user = user
+        self.ip = ip
+
+    def __enter__(self):
+        add_allow_record(self.rest_url, self.username, self.user, self.ip)
+        logger.info("add allow record user %s, ip %s", self.user, self.ip)
+        return self
+
+    def __exit__(self, type, value, traceback):
+        try:
+            delete_allow_record(self.rest_url, self.username, self.user)
+            logger.info("rollback allow record for user %s", self.user)
+        except Exception:
+            logger.exception("failed to rollback allow record for user %s",
+                             self.user)
+
+
 def block_until_state(rest_url, jid, not_in, states, timeout=300):
     start = datetime.datetime.now()
 

src/RestAPI/dlwsrestapi.py

@@ -1157,6 +1157,45 @@ def post(self):
         return JobRestAPIUtils.set_repair_message(username, job_id, payload)
 
 
+@api.resource("/AllowRecord")
+class AllowRecord(Resource):
+    def __init__(self):
+        self.get_parser = reqparse.RequestParser()
+        self.get_parser.add_argument("userName", required=True)
+        self.get_parser.add_argument("user", required=True)
+
+        self.post_parser = reqparse.RequestParser()
+        self.post_parser.add_argument("userName", required=True)
+        self.post_parser.add_argument("user", required=True)
+        self.post_parser.add_argument("ip", required=True)
+
+        self.delete_parser = reqparse.RequestParser()
+        self.delete_parser.add_argument("userName", required=True)
+        self.delete_parser.add_argument("user", required=True)
+
+    def get(self):
+        args = self.get_parser.parse_args()
+        username = args.get("userName")
+        user = args.get("user")
+        resp, code = JobRestAPIUtils.get_allow_record(username, user)
+        return resp, code
+
+    def post(self):
+        args = self.post_parser.parse_args()
+        username = args.get("userName")
+        user = args.get("user")
+        ip = args.get("ip")
+        resp, code = JobRestAPIUtils.add_allow_record(username, user, ip)
+        return resp, code
+
+    def delete(self):
+        args = self.delete_parser.parse_args()
+        username = args.get("userName")
+        user = args.get("user")
+        resp, code = JobRestAPIUtils.delete_allow_record(username, user)
+        return resp, code
+
+
 @app.route("/metrics")
 def metrics():
     return Response(prometheus_client.generate_latest(),

src/docker-images/allowlist-manager/Dockerfile

@@ -0,0 +1,8 @@
+FROM python:3.7
+
+RUN pip3 install requests python-dateutil pyyaml pytz
+RUN curl -sL https://aka.ms/InstallAzureCLIDeb | bash
+
+WORKDIR /allowlist-manager
+
+ADD main.py /allowlist-manager/

src/docker-images/allowlist-manager/allow.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+
+import argparse
+
+from main import RestUtil
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--user",
+                        "-u",
+                        required=True)
+    parser.add_argument("--ip",
+                        "-i",
+                        required=True)
+    parser.add_argument("--rest",
+                        "-r",
+                        required=True)
+    args = parser.parse_args()
+
+    rest_util = RestUtil({"rest_url": args.rest})
+    rest_util.add_allow_record(args.user, args.ip)
+

src/docker-images/allowlist-manager/deny.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+
+import argparse
+
+from main import RestUtil
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--user",
+                        "-u",
+                        required=True)
+    parser.add_argument("--rest",
+                        "-r",
+                        required=True)
+    args = parser.parse_args()
+
+    rest_util = RestUtil({"rest_url": args.rest})
+    rest_util.delete_allow_record(args.user)
+