#!/usr/bin/ruby -w

#
# ntpstat.rb v1.0 - 2007/04
#
# Short: ntpstat.rb displays the status report of a local NTP server.
#
# Long: ntpstat is a program written in C by G.Richard Keech 
#   < rkeech (at) redhat com > that is currently used in most RedHat 
#   based Linux distributions to obtain a succint status report of 
#   the local NTP server.
#	ntpstat.rb is a rewrite in Ruby of ntpstat v0.2 (2001-06-22),
#	which is copyrighted by its author and distributed under the GPLv2.
#
#	the ntp server is contacted via an NTP control message mode 6 that
#	requests a list of state variables. these variables are then parsed
#	and an output is produced to quickly summarise the status of the server.
#
#	since ntp is a time keeping protocol, the main interest of this program
#	is to find out if the server is synchronised with its sources and
#	how accurate is the current time.
#
# I wrote this for fun and to be used from other ruby programs without
# having to deal with more external binary dependencies.
#
# Copyright (C) 2007 by Pedro Venda < pjvenda (at) pjvenda org >
#
# Distributed under the terms of the GNU Public License Agreement (GPLv2)
# http://www.fsf.org/licensing/licenses/gpl.html or
# http://www.fsf.org/licensing/licenses/gpl.txt
#
# Credit
# ======
#
# Main credit goes to the author of ntpstat G.Richard Keech, who, by
# choosing to distribute his program freely under an open source license
# allowed me to easily port the code into another language.
#
# Also, the NTP RFC (1305) was obviously useful for being a reference 
# of the NTP protocol.
#
# Usage
# =====
#
# this script runs from the command line and outputs into stdout:
#
# $ ruby ntpstat.rb
# unsynchronised
#   time server re-starting
#    polling server every 64 s
#
# $ ruby ntpstat.rb
# synchronised to NTP server (10.0.14.1) at stratum 4
#   time correct to within 99 ms
#    polling server every 1024 s
#
# alternatively it can be included as a Ruby class 'Ntpstat':
#
# require 'ntpstat'
# cm=Ntpstat.new
# cm.get
# cm.parse
# cm.to_s
#
# Changelog
# =========
#
# v1.0 - 2007/04/04
# - fully functional port of ntpstat-0.2 to Ruby
#

require 'socket'

class Ntpstat
	# a lot of important constants
	@@NTP_ADDR="127.0.0.1"
	@@NTP_PORT=123
	@@NTP_CM_OP_READVAR="\x02" # 5 last bits
	@@NTP_CM_REQUEST_FLAG1="\x16"	# LI: 00
									# Version: 010 (reserved)
									# Mode: 110 NTP Control Message 6
	@@NTP_CM_REQUEST_FLAG2="\x02"   # Response bit: 0
									# Error bit: 0
									# More bit: 0
									# Opcode: 00010 (READVAR)
	@@NTP_CLK_SRC_NAME=["unspecified","atomic clock","VLF radio","HF radio","UHF radio","local net","NTP server","UDP/TIME","wristwatch","modem"]
	@@MAX_UDP_SIZE=1024
	@@NTP_CM_PAYLOAD_SIZE=468
	@@NTP_CM_AUTHENTICATOR_SIZE=96

	# accessible output is info hash and error string
	attr_reader :info, :error
	# configuration options - server address and port
	attr_accessor :addr, :port

	def initialize(addr=nil,port=nil)
		# info hash will contain parsed NTP CM response
		@info={}
		# error string if anything went wrong
		@error=nil
		# ntp server address
		if addr != nil
			# given by user
			@addr=addr
		else
			# default - localhost
			@addr=@@NTP_ADDR
		end
		if port != nil
			# given by user
			@port=port
		else
			# default - 123
			@port=@@NTP_PORT
		end
		# raw UDP packet with server response
		@reply=nil
	end

	#
	# requests data by sending an NTP control mode 6 message
	# to the ntp server.
	# receives a response and stores it in @reply
	#
	def get
		# NTP Control Mode 6 Message Fields (hex)
		#
		#    Flags: 0x16
		#        00.. .... = Leap Indicator: no warning (0)
		#        ..01 0... = Version number: reserved (2)
		#        .... .110 = Mode: reserved for NTP control message (6)
		#    Flags 2: 0x02
		#        0... .... = Response bit: Request (0)
		#        .0.. .... = Error bit: 0
		#        ..0. .... = More bit: 0
		#        ...0 0010 = Opcode: readvar (2)
		cmd=@@NTP_CM_REQUEST_FLAG1+@@NTP_CM_REQUEST_FLAG2
		seqval="\x01"  # sequence value - anything
		status="\x00"  # status values - anything
		associd="\x00" #
		offset="\x00"  #
		count="\x00"   #
	
		# NTP Control Mode 6 Message PACK sequence
		pack_string="a2a2a2a2a2a2c*"
	
		begin
			# open client socket
			csock=UDPSocket.open
		rescue # rescue all exceptions...
			@error="unable to open socket"
			return nil
		end
	
		begin
			# connect to NTP_ADDR (localhost) and NTP_PORT (123)
			csock.connect(addr,port)
		rescue # rescue all exceptions...
			@error="unable to connect to socket"
			return nil
		end
	
		# build message array and pack it into a stream of bytes - the udp packet
		ntpmsg_array=[cmd,seqval,status,associd,offset,count]+[0]*(@@NTP_CM_PAYLOAD_SIZE+@@NTP_CM_AUTHENTICATOR_SIZE)
		ntpmsg=ntpmsg_array.pack(pack_string)
		begin
			# send request
			csock.send(ntpmsg,0)
		rescue # rescue all exceptions...
			@error="unable to send command to NTP port"
			return nil
		end

		begin
			# try to obtain a response from the same source (NTP_ADDR:NTP_PORT)
			reply=csock.recvfrom(@@MAX_UDP_SIZE)
		rescue # rescue all exceptions...
			@error="Unable to talk to NTP daemon. Is it running?"
			return nil
		end
		@reply=reply[0]
	end
	
	#
	# parses the received message from the ntp server
	# results in either a hash @info with relevant parsed information
	# or an @error string if anything went wrong.
	#
	def parse
		# unpack string - to extract different fields from response packet
		unpack_string="B16H4B16H4H4H4a*"
		if @reply==nil
			return nil
		end

		# unpack response packet according to unpack_string instructions
		ntpreply=@reply.unpack(unpack_string)
	
		# do a little manipulation on the reply
		# build a hash with the returned results
		ntpreply_hash={}

		# calculate clock source code
		# status byte 2, bits 3-8
		ntpreply_hash[:clock_source]=ntpreply[2][2,6].to_i(2)

		# recover command bytes for potential future access
		ntpreply_hash[:command]=ntpreply[0].split(//)
		ntpreply_hash[:sequence_number]=ntpreply[1]
		ntpreply_hash[:status_code]=ntpreply[2].split(//)
		ntpreply_hash[:association_id]=ntpreply[3]
		ntpreply_hash[:offset_size]=ntpreply[4]
		ntpreply_hash[:count_number]=ntpreply[5]
		# recover several flags for potential future access
		if ntpreply_hash[:command][8]=="1"
			ntpreply_hash[:response_bit]=true
		else
			ntpreply_hash[:response_bit]=false
		end
		if ntpreply_hash[:command][9]=="1"
			ntpreply_hash[:error_bit]=true
		else
			ntpreply_hash[:error_bit]=false
		end
		if ntpreply_hash[:command][10]=="1"
			ntpreply_hash[:more_bit]=true
		else
			ntpreply_hash[:more_bit]=false
		end

		# the following field is a big string with a lot of variables
		# kindly extracted from the status fields by the server
		#
		# remove any \n or \n left and split fields by ',' to put elements on
		# the hash
		ntpreply=ntpreply[6].gsub(/[\n\r]/,"").split(/, ?/)
		ntpreply.flatten!
		ntpreply.each do |item|
			ar=item.split("=")
			ntpreply_hash[ar[0].to_sym] = ar[1]
		end

		# change data type of some known entries
		# Integers
		ntpreply_hash[:stratum]=ntpreply_hash[:stratum].to_i
		ntpreply_hash[:poll]=ntpreply_hash[:poll].to_i
		ntpreply_hash[:state]=ntpreply_hash[:state].to_i
		ntpreply_hash[:leap]=ntpreply_hash[:leap].to_i
		ntpreply_hash[:precision]=ntpreply_hash[:precision].to_i
		ntpreply_hash[:peer]=ntpreply_hash[:peer].to_i
		# Floats
		ntpreply_hash[:noise]=ntpreply_hash[:noise].to_f
		ntpreply_hash[:jitter]=ntpreply_hash[:jitter].to_f
		ntpreply_hash[:stability]=ntpreply_hash[:stability].to_f
		ntpreply_hash[:offset]=ntpreply_hash[:offset].to_f
		ntpreply_hash[:frequency]=ntpreply_hash[:frequency].to_f
		ntpreply_hash[:rootdispersion]=ntpreply_hash[:rootdispersion].to_f
		ntpreply_hash[:rootdelay]=ntpreply_hash[:rootdelay].to_f
		# Strings
		ntpreply_hash[:version].gsub!("\"","")
		ntpreply_hash[:system].gsub!("\"","")
		ntpreply_hash[:processor].gsub!("\"","")

		#
		# check for response message validity:
		#
		# - first command byte must be as sent except for the leap indicator
		# 	command byte 1 bits 3-8
		# - second command byte must have response=1, error=0 and more=0
		#   command byte 2 bits 1, 2 and 3
		if ntpreply_hash[:command][2,6].join != @@NTP_CM_REQUEST_FLAG1.unpack("B8")[0][2,6]
			@error="return data appears to be invalid based on status word"
			return nil
		end
		if not ntpreply_hash[:response_bit]
			@error="response bit not set in reply"
			return nil
		end
		if ntpreply_hash[:error_bit]
			@error="error bit is set in reply"
			return nil
		end
		if ntpreply_hash[:more_bit]
			@error="More bit unexpected in reply"
			return nil
		end

		@info=ntpreply_hash
	end
	
	#
	# produces an output string from the internal variables
	# @info or @error (if any)
	#
	def to_s
		# @info is not filled if anything went wrong
		if @info=={}
			if @error != nil
				# if anything went wrong and the program got it,
				# then an error message is available in @error
				return @error
			else
				return ""
			end
		end

		ret=""
		# leap indicator (LI) shows unsynchronisation if both bits are 1
		if @info[:leap]==3
			ret+="unsynchronised\n"
			# if state is 1, then the server is restarting
			if @info[:state]==1
				ret+="  time server re-starting\n"
			end
		else
			# clock is synchronised
			ret+="synchronised to "
			# clock_source is reference type (GPS/ATOMIC/NTP/whatever)
			if @info[:clock_source] < 10
				ret+=@@NTP_CLK_SRC_NAME[(@info[:clock_source])]
				if @info[:clock_source]==6 && @info[:refid]!=nil
					if @info[:refid].length < 15 && @info[:refid].length>0
						# refid is main server address
						ret+=" ("+@info[:refid]+")"
					else
						ret+=" <IP unreadable>"
					end
				end
			else
				ret+=" unknown source"
			end
			# local synchronisation stratum
			if @info[:stratum]
				ret+=" at stratum "+@info[:stratum].to_s+"\n"
			else
				ret+=", stratum unknown\n"
			end
			# root dispersion is synchronisation error range
			if @info[:rootdispersion]
				ret+="   time correct to within "+@info[:rootdispersion].truncate.to_s+" ms\n"
			else
				ret+="   accuracy unknown\n"
			end
		end
		# polling interval is 2^poll
		if @info[:poll]
			ret+="   polling server every "+(2**@info[:poll]).to_s+" s\n"
		else
			ret+="   polling interval unknown\n"
		end
		ret
	end

end

# see if we were called directly
if $0 == __FILE__
	cm=Ntpstat.new
	# request information and receive response
	cm.get
	# parse response string
	cm.parse
	# print summarised display
	puts cm.to_s
end