#!/usr/bin/python3

# Readme
# ======
# (parts copy/pasted from redd.it/fjgit)
# It works like this:
#
#     $ tgrep 8:42:04
#     [log lines with that precise timestamp]
#     $ tgrep 10:01
#     [log lines with timestamps between 10:01:00 and 10:01:59]
#     $ tgrep 23:59-0:03
#     [log lines between 23:59:00 and 0:03:59]
#
# By default it uses /logs/haproxy.log as the input file, but you can specify 
# an alternate filename by appending it to the command line. It also works if 
# you prepend it, because who has time to remember the order of arguments for 
# every little dumb script?
#
# Most importantly, tgrep is fast, because it doesn't look at every line in 
# the file. It jumps around, checking timestamps and doing a binary search 
# until it finds the range you're looking for.
#
# This program assumes all lines in the log file will matches the following
# pattern, which is not an unreasonable assumption, but expect an ugly error 
# message if a line doesn't match.
# Log pattern: \w+\s+\d+\s+\d+:\d+:\d+

# Out of Order Time Stamps
# ========================
# In the 'real world' the time stamps aren't going to always increase, but
# the log will be in mostly sorted order. max_displacement is included to 
# increase the range read near by a fixed number of bytes which enables finding 
# results that are out of order. The actual value of the parameter would need to
# be tuned based on real world logs to make sure no data is missed without
# making performance worse.

# Special Cases
# =============
# The log file spans at most two days (crosses midnight at most once), so it
# can be broken into two parts: the first day, and possibly a second.
#
# Start time and end time are provided on the command line (end time defaults
# to the start time). If the end time < start time, the span crosses midnight,
# in which case the program checks the time ranges (start, midnight) and
# (midnight, end) for each day. Otherwise each day is just checked for entries
# in the range (start, end).
#
# If the time range matches on both days, the results from both days will be
# shown (maintaining the order in the log)

# Performance
# ===========
#     $ tgrep 6:54:03-6:54:05 > output.txt
# (max_displacement = 0)
# Input: 72.9 GB file (2^14 entries per second over 6:52 - 7:13 the next day)
# Output: 5.7 MB file, 98304 entries
# Made 3548 system calls, ran in 0.8 seconds (7200rpm SATA300 WDC drive)
# The max resident size in memory was under 27 MB.
# Setting max_displacement to 1 MB increases the run time to 1.3 seconds
# Setting max_displacement to 6 MB increases the run time to 4 seconds
#
#     $ tgrep 8:23:54 > output.txt 
# took 0.25 seconds and produced one-sixth the output, signaling that the
# amount of data being written out is likely the bottle neck.
# Setting max_displacement to 1 MB increases the run time to 0.55 seconds
# Setting max_displacement to 6 MB increases the run time to 2.6 seconds
#
# It could be made a bit faster by not using regex's and datetime.time if one 
# second is just too long to wait.
#
# Big-O average is O(lg(n)) with respect to the size of the file, but the
# pathological case where the last line is more than half the file would make
# it O(n).

import re
import io
import sys
from datetime import time

# Default logs, open the first of these that exists if a file isn't given
default_logs = ['/log/haproxy.log']

# Out of order parameter, number of extra byte to check before and after the
# actual found range
max_displacement =  1 * (1024**2)

# Regex's to match beginning of lines in the file and command line arguments
log_pattern = re.compile(r"(?P<date>(?P<month>\w+)\s+(?P<day>\d+))\s+(?P<time>(?P<hour>\d+):(?P<min>\d+):(?P<sec>\d+))")
ts_pattern = re.compile(r"(?P<hour>\d+)(:(?P<min>\d+)(:(?P<sec>\d+))?)?")

# Convert a match to a time stamp, optionally rounding up instead of down
def group_time(m, down = True):
  a,b,c = m.group('hour'), m.group('min'), m.group('sec')
  if down:
    return time(a!=None and int(a) or 0, b!=None and int(b) or 0, c!=None and int(c) or 0, 0)
  else:
    return time(a==None and 23 or int(a), b==None and 59 or int(b), c==None and 59 or int(c), 999999)

class ParserException(Exception):
  pass

# Parse a string as a time or time range
def parse_range(r):
  times = r.split('-')
  if len(times) == 1:
    times.append(times[0])
  times = [ts_pattern.match(x) for x in times]
  if len(times) != 2 or not all(times):
    raise ParserException()
  start = group_time(times[0])
  end = group_time(times[1], False)
  return start,end

# Find the seek position of the first line such that f(line) is true.
# f(line) should be interpreted as line >= the desired line.
# Optionally takes bounds for the min and max seek position to start lines
def first_line(source, f, low=-1, high=-1):
  if low == -1: # Default bounds
    low = 0
    high = source.seek(0,io.SEEK_END)
  # Binary search
  while True:
    mid = (low+high)//2
    source.seek(mid)
    # Read incomplete line and next full line
    advance = source.readline()
    line = source.readline()
    # Adjust mid to beginning of the full line
    mid += len(advance)
    if mid == high:
      break
    if f(line):
      high = mid
    else:
      low = mid
  source.seek(low)
  # The binary search ensures that if the desired line exists, it starts
  # between low and high, so search linearly returning the first match.
  # This loop is only executed a few times if the lines are roughly similar
  # length since there is no newline between mid and high.
  while low <= high:
    line = source.readline()
    if len(line) == 0:
      break
    if f(line):
      return low
    low += len(line)
  return -1

def print_lines(source, start, end):
  # Find the point where midnight is crossed, if any.
  # This only works because it midnight is crossed at most once.
  file_size = source.seek(0, io.SEEK_END)
  source.seek(0)
  first_day = log_pattern.match(source.readline()).group('date')
  i = first_line(source, lambda x: log_pattern.match(x).group('date') != first_day)
  if i < 0:
    days = [(0, file_size)]
  else:
    days = [(0, i-1), (i, file_size)]

  # The ranges of time to search, if it crosses midnight consider it to be
  # two times: start-midnight, midnight-end
  if end < start:
    time_ranges = [(time.min, end), (start, time.max)]
  else:
    time_ranges = [(start, end)]

  read_upto = 0 # Mark the next earliest line to read
  # Check all days in each time range
  for day_start, day_end in days:
    for time_start, time_end in time_ranges:
      # Find where the first match (maybe) is, if time_start == time.min, we
      # could just let i = day_start since time.min <= all other times, but
      # there may not be any matches between midnight and time_end
      s_start = first_line(source, lambda x: group_time(log_pattern.match(x)) >= time_start, day_start, day_end)
      if s_start < max_displacement:
        s_start = max_displacement
      # i is the position to start reading at
      i = s_start - max_displacement # i >= 0
      source.seek(i)
      i += len(source.readline()) # move to start of first full line
      if s_start < i: # if max_displacement < len(line) or at start of file
        i = s_start
        source.seek(i)
      if i < read_upto: # Fastfoward if we've already been here
        i = read_upto
        source.seek(i)
      # j is the max seek position to read until
      j = day_end + max_displacement
      if j > file_size:
        j = file_size
      # Print lines until the end of the day or time range
      while i < j:
        line = source.readline()
        i += len(line)
        read_upto = i
        t = group_time(log_pattern.match(line), False)
        if t <= time_end and t >= time_start:
          print(line,end='')
        elif t > time_end and i + max_displacement < j:
          # Since this record is past the desired range
          # only read for another max_displacement bytes
          j = i + max_displacement

def main():
  start, end, log = None, None, None
  args = sys.argv[1:]

  # Find an argument which can be interpreted as a timestamp or range
  for arg in args:
    try:
      start,end = parse_range(arg)
      break
    except ParserException:
      pass

  if not start or not end:
    print("Error: No valid timestamp or range found.", file=sys.stderr)
    sys.exit()
    
  # Find an argument which is a file that can be opened, default to /log/haproxy.log
  for arg in args + default_logs:
    try:
      log = open(arg,'r')
      break
    except IOError:
      pass

  if not log:
    print("Error: No input file provided.", file=sys.stderr)
    sys.exit()

  # Get 'er done.
  print_lines(log, start, end)

if __name__ == '__main__':
  main()