set uid and gid

src/agentx-ntpsec.py

@@ -4,6 +4,8 @@ import time
 import argparse
 import os
 import sys
+import pwd
+import grp
 import logging
 import logging.handlers
 
@@ -64,6 +66,23 @@ def daemonize(pid_filename):
     logger.removeHandler(stdout_handler)
 
 
+def set_user_group(user, group):
+    if group:
+        try:
+            gid = grp.getgrnam(group).gr_gid
+        except KeyError:
+            logger.error(f"Group {group} does not exist")
+            sys.exit(1)
+        os.setgid(gid)
+    if user:
+        try:
+            uid = pwd.getpwnam(user).pw_uid
+        except KeyError:
+            logger.error(f"user {user} does not exist")
+            sys.exit(1)
+        os.setuid(uid)
+
+
 if __name__ == '__main__':
     logging.basicConfig(
             level=logging.DEBUG, 
@@ -96,11 +115,36 @@ if __name__ == '__main__':
                         help=f"pid-file when running as daemon, default is {pid_filename}",
                         required=False,
                         default=pid_filename)
+    parser.add_argument('-u', '--user',
+                        help="Set uid of process",
+                        required=False,
+                        default='')
+    parser.add_argument('-g', '--group',
+                        help="Set gid of process",
+                        required=False,
+                        default='')
 
     args = parser.parse_args()
     
     if args.daemonize:
         daemonize(pid_filename)
+    set_user_group(args.user, args.group)
+
+    if args.group:
+        try:
+            gid = grp.getgrnam(args.group).gr_gid
+        except KeyError:
+            logger.error(f"Group {args.group} does not exist")
+            sys.exit(1)
+        os.setgid(gid)
+    if args.user:
+        try:
+            uid = pwd.getpwnam(args.user).pw_uid
+        except KeyError:
+            logger.error(f"user {args.user} does not exist")
+            sys.exit(1)
+        os.setuid(uid)
+        
 
     ntpserver = args.ntpserver
     period = args.period