Some optimizations, systemd service, native socket, and rpm spec file

motd-generator.spec

@@ -0,0 +1,96 @@
+Name:		motd-generator
+Version:	1.0.1
+%define build_timestamp %{lua: print(os.date("%Y%m%d"))}
+Release:	%{build_timestamp}
+Summary:	Custom message of the day (MOTD) designed to be as practical and informative as possible.
+BuildArch:	noarch
+
+License:	Unknown
+URL:		https://github.com/gonoph/motd-generator
+Source0:	https://github.com/gonoph/motd-generator/archive/systemd.tar.gz
+
+Requires:	systemd
+Requires:	python
+Requires:	python(abi) = 2.7
+
+Provides:	motd
+
+%{?systemd_requires}
+BuildRequires:	systemd
+Buildrequires:	python
+
+%description
+This is a custom message of the day (MOTD) designed to be as practical and
+informative as possible. The truth is, no one actually reads the MOTD. As such,
+the MOTD should contain useful, yet minimal, information about the host system
+such that a quick glance at it when logging in may actually be worth a person's
+precious time. This way, any potential issues are noticed and not naively
+ignored. This MOTD generator scripts has the ability to output text in color.
+Using this feature, potential issues can be highlighted for easy
+identification.
+
+
+%prep
+%autosetup -n %{name}-systemd
+
+# %build
+
+
+%install
+DESTDIR=%{buildroot}
+DIR_BIN=${DESTDIR}%{_bindir}
+DIR_STATE=${DESTDIR}%{_sharedstatedir}/%{name}
+DIR_SYSTEMD=${DESTDIR}%{_exec_prefix}/lib/systemd/system
+DIR_PRESETS=${DESTDIR}%{_exec_prefix}/lib/systemd/system-preset/
+DIR_PROFILE=${DESTDIR}%{_sysconfdir}/profile.d
+mkdir -p $DIR_BIN
+mkdir -p $DIR_STATE
+mkdir -p $DIR_SYSTEMD
+mkdir -p $DIR_PRESETS
+mkdir -p $DIR_PROFILE
+install -m 755 -t $DIR_BIN motd_gen.py motd_stat.py
+install -m 644 -t $DIR_SYSTEMD motd_stat.service
+echo "enable motd_stat.service" > $DIR_PRESETS/50-motd_stat.preset
+cat << EOF > $DIR_PROFILE/motd_gen.sh
+if [ -e %{_bindir}/motd_gen.py ]; then
+    if [ -x %{_bindir}/dircolors ]; then
+        %{_bindir}/motd_gen.py --color --warn --border
+    else
+        %{_bindir}/motd_gen.py --border
+    fi
+fi
+EOF
+
+%pre
+DIR_STATE=${DESTDIR}%{_sharedstatedir}/%{name}
+U=$(getent passwd motd)
+test "x$U" == x && useradd -r -d $DIR_STATE -M -U motd || :
+
+%post
+%systemd_post motd_stat.service
+
+%preun
+%systemd_preun motd_stat.service
+
+%postun
+%systemd_postun_with_restart motd_stat.service
+if [ $1 -gt 0 ] ; then
+  exit 0
+fi
+userdel motd || :
+
+%files
+%defattr(644, motd, motd, 755)
+%attr(755,-,-) %{_bindir}/motd_gen.py
+%attr(755,-,-) %{_bindir}/motd_stat.py
+%{_sharedstatedir}/%{name}
+%{_exec_prefix}/lib/systemd/system/motd_stat.service
+%{_exec_prefix}/lib/systemd/system-preset/50-motd_stat.preset
+%{_sysconfdir}/profile.d/motd_gen.sh
+%doc README.md
+%doc motd_stat
+%dir %{_sharedstatedir}/%{name}
+
+%changelog
+* Fri Sep 14 2018 Billy Holmes <[email protected]> - 1.0.0-1
+- Initial release

motd_gen.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python
+# vim: ts=4 sw=4 ai expandtab
 
 # Written in 2012 by Joe Tsai <[email protected]>
 #
@@ -27,6 +28,9 @@
 import getpass
 import optparse
 import datetime
+import socket
+from StringIO import StringIO as SIO
+import traceback
 
 
 ################################################################################
@@ -91,6 +95,7 @@
 CACHE_FREE = True # Is disk cache considered free memory or not?
 FULL_HOSTNAME = False # Use the full FQDN hostname
 STAT_PORT = 4004 # Port for the motd_netstat daemon
+STAT_HOST = 'localhost' # Host for the motd_netstat daemon
 NETTRAF_DEVICE = 'eth0' # The network device to monitor
 NETTRAF_INTERVAL = 600 # Time length in seconds to average the bandwidth over
 NETTRAF_WEIGHT = 1.0 # Perform linear moving average weight
@@ -107,6 +112,34 @@
 ############################### Helper functions ###############################
 ################################################################################
 
+def send_to_port(host, port, payload):
+    for res in socket.getaddrinfo(host, port, socket.AF_INET, socket.SOCK_STREAM):
+        af, socktype, proto, canonname, sa = res
+        received = SIO()
+        s = socket.socket(af, socktype, proto)
+        try:
+            s.connect(sa)
+            s.send(payload)
+            for tmp in s.recv(1024):
+                received.write(tmp)
+            return received.getvalue()
+        finally:
+            s.close()
+            received.close()
+
+def extract_error():
+    """Extracts meaning error from a caught exception"""
+    info = sys.exc_info()
+    tb = traceback.extract_tb(info[2])[0]
+    _file = tb[0]
+    _line = tb[1]
+    return "%s(%s) at %s line %s" % (
+        colorize('exception',CYAN0),
+        colorize(str(info[1]), YELLOW0),
+        colorize(_file, YELLOW1),
+        colorize(_line, YELLOW1)
+        )
+
 def exec_cmd(cmd):
     """Execute a command and retrieve standard output"""
     with os.popen(cmd + ' 2> /dev/null', 'r') as cmd_call:
@@ -214,11 +247,22 @@ def display_lower_border():
     """Display the lower border"""
     display_border(LOWER_HALF_BLOCK)
 
+def _strip_list(l):
+    return map(lambda x: x.strip('\r\n\t "'), l)
 
 def display_welcome():
     """Display the welcome message"""
-    os_issue = ''.join(read_file('/etc/issue')).strip()
-    os_name = re.sub(r'\\[a-zA-Z]', '', os_issue).strip()
+    os_release = None
+    os_name = None
+    try:
+        _filter = filter(lambda x: x, _strip_list(read_file('/etc/os-release') ) )
+        os_release = dict(map( lambda x: _strip_list(x.split('=') ), _filter) )
+        os_name = "%s %s" % (os_release['NAME'], colorize(os_release['VERSION'], GREEN1))
+    except:
+        pass
+    if not os_release:
+        os_issue = ''.join(read_file('/etc/issue')).strip()
+        os_name = re.sub(r'\\[a-zA-Z]', '', os_issue).strip()
     cmd = 'hostname -f' if FULL_HOSTNAME else 'hostname'
     host_name = ''.join(exec_cmd(cmd)).strip()
     values = colorize(host_name, TEXT_PRIMARY), colorize(os_name, TEXT_PRIMARY)
@@ -239,6 +283,14 @@ def display_info():
         key = (key + ':').ljust(max_length + 4, ' ')
         print " %s%s" % (colorize(key, TEXT_SECONDARY), value)
 
+def display_warnings():
+    """Display any warnings from the WARN_LOG"""
+    global WARN_LOG
+    global opts
+    if not opts.warn:
+        return
+    for warning in WARN_LOG:
+        print " - %s %s" % (colorize('WARNING:', WARNING), colorize(warning, RESET))
 
 ################################################################################
 ################################ Options parser ################################
@@ -325,6 +377,8 @@ def display_info():
 # Generate info list
 result_list = []
 
+WARN_LOG = []
+
 # Get last login
 try:
     login_host = exec_cmd('last -n 2 -w -F $USER')
@@ -429,9 +483,8 @@ def display_info():
                 'weight':   CPUUTIL_WEIGHT,
             }
         }
-        query = shell_escape(json.dumps(query))
-        command = 'echo %s | netcat localhost %s' % (query,STAT_PORT)
-        util_usage = ''.join(exec_cmd(command)).strip()
+        query = json.dumps(query)
+        util_usage = send_to_port(STAT_HOST, STAT_PORT, query).strip()
 
         # Load the JSON data
         data = json.loads(util_usage)
@@ -446,8 +499,8 @@ def display_info():
     values = tuple(utils_text)
     message = "%s (1 minute) - %s (5 minutes) - %s (15 minutes)" % values
     info_list.append(('CPU utilization', message))
-except:
-    pass
+except Exception as ex:
+    WARN_LOG.append(extract_error())
 
 # Get CPU load
 try:
@@ -519,9 +572,8 @@ def display_info():
             'weight':   NETTRAF_WEIGHT,
         }
     }
-    query = shell_escape(json.dumps(query))
-    command = 'echo %s | netcat localhost %s' % (query, STAT_PORT)
-    net_usage = ''.join(exec_cmd(command)).strip()
+    query = json.dumps(query)
+    net_usage = send_to_port(STAT_HOST, STAT_PORT, query).strip()
 
     # Load the JSON data
     data = json.loads(net_usage)
@@ -534,8 +586,8 @@ def display_info():
     values = total_text,units(rx_avg, 'B/s'), units(tx_avg, 'B/s')
     message = "%s - %s down, %s up" % values
     info_list.append(('Network traffic', message))
-except:
-    pass
+except Exception as ex:
+    WARN_LOG.append(extract_error())
 
 # Get processes
 try:
@@ -555,4 +607,5 @@ def display_info():
 display_welcome()
 display_logo()
 display_info()
+display_warnings()
 display_lower_border()

motd_stat.service

@@ -0,0 +1,18 @@
+[Unit]
+Description=dsnet motd stat service
+After=sshd.service
+After=getty.target
+
+[Service]
+Type=simple
+Restart=always
+RestartSec=5
+WorkingDirectory=/var/lib/motd-generator
+User=motd
+Group=motd
+ExecStart=/usr/bin/motd_stat.py
+MemoryLimit=128M
+CPUQuota=20%
+
+[Install]
+WantedBy=multi-user.target