#! /usr/bin/python

import sys, os, dircache, string

#-----------------------------------------------------------------------------
# Required Subroutines
#-----------------------------------------------------------------------------

def remove_dups():
#-----------------------------------------------------------------------------
# Display duplicate package files and prompt for removal of the last shown
# by sort order, the most recent by date, or none
#-----------------------------------------------------------------------------
  global ms, me, df, flag

  high = ms
  for n in range(ms+1, me+1):
    if dd[n] >= dd[high]:
      high = n

  print ""
  for n in range(ms, me+1):
    if n == high:
      print "*",
    else:
      print " ",
    print df[n]
    
  if (flag == "P"):
    print "\nKeep only " + df[me] + " (y/N/d/q) ? ",
    yn = sys.stdin.readline()
  elif (flag == "S"):
    yn = "Y\n"
  elif (flag == "D"):
    yn = "D\n"
    
  if yn == "\n":
    yn = "N\n"
  yn = string.upper(yn[:1])

  if yn == "Y":
    for i in range(ms, me):
      print "\bRemoving: " + basepath + df[i]
      os.remove(basepath + df[i])
    return -1

  if yn == "D":
    for i in range(ms, me+1):
      if i != high:
        print "\bRemoving: " + basepath + df[i]
        os.remove(basepath + df[i])
    return -1

  if yn == "Q":
    sys.exit(0)

  return 0

def prompt_dups():
#-----------------------------------------------------------------------------
# Display each duplicate package file and prompt to keep, or remove
#-----------------------------------------------------------------------------
  global ms, me, df, basepath

  for n in range(ms, me+1):
    print "\bKeep:", df[n], " (Y/n/q) ? ",

    yn = sys.stdin.readline()
    if yn == "\n":
      yn = "Y\n"
    yn = string.upper(yn[:1])

    if yn == "Q":
      sys.exit(0)

    if yn == "N":
      print "\bRemoving: " + basepath + df[n]
      os.remove(basepath + df[n])

def is_installed(name):
#-----------------------------------------------------------------------------
# Check if a package is currently installed.
# Return 0 if uninstalled and 1 if installed. 
#-----------------------------------------------------------------------------
  global di

  installed = 0

  for n in range(len(di)):
    if (di[n] == name):
      installed = 1
      break

  return installed

def is_date(ds):
#-----------------------------------------------------------------------------
# Check if a datestring is valid
#-----------------------------------------------------------------------------
  if len(ds) != 10:
    return 0

  if (ds[4] != "-") or (ds[7] != "-"):
    return 0

  n = int(ds[0:4])
  if (n < 1970) or (n > cyear):
    return 0

  n = int(ds[5:7])
  if (n < 1) or (n > 12):
    return 0

  n = int(ds[8:])
  if (n < 1) or (n > 31):
    return 0  
  
  return -1
    
#-----------------------------------------------------------------------------
# End of Subroutines
#
# Main program starts here
#-----------------------------------------------------------------------------

print "CleanApt v0.4.0 - an apt cache cleaner\n"

#-----------------------------------------------------------------------------
# Check for -h or --help argument passed to program and display help
# if requested.
#-----------------------------------------------------------------------------

args = len(sys.argv) - 1

if args > 1:
  print "\nUsage: " + sys.argv[0][2:] + " [-h | --help | -r | --remove | -d | --date\n"
  sys.exit(0)
  
flag = "P"

if (args == 1):
  if ((sys.argv[1] == "-h") or (sys.argv[1] == "--help")):
    print "Remove duplicate, or unnecessary packages from the apt archives.\n"

    print "After displaying some archives statistics, the following prompts will be"
    print "presented:\n"

    print "Delete uninstalled packages from archive (y/N/p/q)? "
    print "   y = Delete all uninstalled package files from the archives"
    print "       without prompting."
    print "   N = Do not delete any package files based on installation status."
    print "   p = Prompt to keep each uninstalled package file in the archive."
    print "   q = Quit program without deleting anything.\n"

    print "If prompting was requested:\n"

    print "Keep: filename.deb (Y/n/q) ? "
    print "   Y = Keep the package file in the archives (do not delete)."
    print "   n = Delete the package file from the archives."
    print "   q = Quit program without deleting the package file.\n\n"

    print "Press <ENTER> to continue."
    sys.stdin.readline()

    print "If you have used the -s or --sort command line argument then for all packages"
    print "with multiple files in the archive all files will be deleted EXCEPT the file"
    print "which comes last in the system sort order.\n"
    
    print "If you have used the -d or --date command line argument then for all packages"
    print "with multiple files in the archive all files will be deleted EXCEPT the file"
    print "which has the most recent date.  If more than one file matches the most recent"
    print "date then the file with that date which comes last in the system sort order is"
    print "the file which will be kept.\n\n\n\n\n\n\n\n\n\n\n\n\n"

    print "Press <ENTER> to continue."
    sys.stdin.readline()

    print "If no command line options have been used, then for each package with multiple"
    print "files in the archive a list of each file will be presented.  The list is"
    print "presented in the system sort order and an asterisk (*) is displayed before the"
    print "file with the most recent date.\n"

    print "For the last package shown the following prompt will be presented.\n" 

    print "Keep only filename.deb (y/N/d/q) ? "
    print "   y = Keep the listed package file (deleting all others for"
    print "       the specified package)."
    print "   N = Prompt to keep each package file."
    print "   d = Keep only the most recent file (as indicated with an asterisk."
    print "   q = Quit program without deleting the package file.\n\n\n\n\n\n\n\n\n"

    print "Press <ENTER> to continue."
    sys.stdin.readline()

    print "If prompting for package deletion was requested:\n"

    print "Keep: filename.deb (Y/n/q) ? "
    print "   Y = Keep the package in the archives (do not delete)."
    print "   n = Delete the package from the archives."
    print "   q = Quit program without deleting the package.\n"

    print "NOTE: All responses are case insensitive and require only a single character."
    print "      Defaults are shown in uppercase.\n"
    print "      Unless command line options override prompting taking the defaults"
    print "      WILL NOT DELETE ANYTHING!\n\n\n\n\n\n\n\n\n\n\n"
    sys.exit(0)
  elif ((sys.argv[1] == "-s") or (sys.argv[1] == "--sort")):
    flag = "S"
  elif ((sys.argv[1] == "-d") or (sys.argv[1] == "--date")):
    flag = "D"
  else:
    print "\nUsage: " + sys.argv[0][2:] + " [-h | --help]\n"
    sys.exit(0)

#-----------------------------------------------------------------------------
# Build arrays with required information from directory and dpkg listings
#-----------------------------------------------------------------------------


basepath = "/var/cache/apt/archives/"

df = []     # Full file name from ls
dn = []     # Package name from ls
dd = []     # Date from ls
di = []     # Installed package name from dpkg -l

#-----------------------------------------------------------------------------
# Month abbreviations and their corresponding numbers
#-----------------------------------------------------------------------------

mn = {"Jan":"01", "Feb":"02", "Mar":"03", "Apr":"04", \
      "May":"05", "Jun":"06", "Jul":"07", "Aug":"08", \
      "Sep":"09", "Oct":"10", "Nov":"11", "Dec":"12"}

#-----------------------------------------------------------------------------
# Get current date
#-----------------------------------------------------------------------------

temp = os.popen("date")
line = temp.readline()
temp.close()

cyear = line[24:-1]
cmmdd = mn[line[4:7]] + line[8:10]

#-----------------------------------------------------------------------------
# Get required information from ls
#-----------------------------------------------------------------------------

temp =  os.popen("ls -l --time-style=long-iso " + basepath)

line=temp.readline()
line=temp.readline()

ns = string.rfind(line, " ") +1

while line:
  if (line[-5:-1] == ".deb"):
    df.append(line[ns:-1])
#    dd.append(line[ns-13:ns-1])
    dd.append(line[ns-17:ns-7])
    us = string.find(line, "_", ns)
    dn.append(line[ns:us])
  line = temp.readline()

temp.close()

dl = len(dn)

#-----------------------------------------------------------------------------
# Fix year if ls supplies time instead of year
#-----------------------------------------------------------------------------

for n in range(dl):
  if is_date(dd[n]) == 0:
    print "Date returned by ls (" + dd[n] + ") is invalid.  Aborting!\n\n"
    sys.exit(0)
  dd[n] = dd[n][0:4] + dd[n][5:7] + dd[n][8:]

#-----------------------------------------------------------------------------
# Get required information from dpkg -l
#-----------------------------------------------------------------------------
  
temp = os.popen("export COLUMNS=300;dpkg -l")
line=temp.readline()
line=temp.readline()
line=temp.readline()
line=temp.readline()
line=temp.readline()
line=temp.readline()
while line:
  if line[1:2] == "i":
    di.append(string.strip(line[4:74]))
  line=temp.readline()
temp.close()

#-----------------------------------------------------------------------------
# Initialize variables
#-----------------------------------------------------------------------------

dup_pkgs = 0
dup_files = 0
non_pkgs = 0

ms = -1
me = -1

for n in range(1,dl):
  if dn[n] == dn[n-1]:
    if ms == -1:
      ms = n-1
  else:
    if ms != -1:
      me = n-1

  if (ms > -1) and (n == dl-1):
    me = dl-1

  inst = 0

  if (ms == -1):
    inst = is_installed(dn[n-1])
    if (inst == 0):
      non_pkgs = non_pkgs + 1
    if (n == dl-1):
      inst = is_installed(dn[n])
      if (inst == 0):
        non_pkgs = non_pkgs + 1
  elif (me > -1):
    dup_pkgs = dup_pkgs + 1
    dup_files = dup_files + me - ms + 1
    inst = is_installed(dn[me])
    if (inst == 0):
      non_pkgs = non_pkgs + 1
    ms = -1
    me = -1
  

tot_pkgs = dl - dup_files + dup_pkgs

#-----------------------------------------------------------------------------
# Convert numeric values to strings for proper alignment
#-----------------------------------------------------------------------------

sti = ("    " + str(len(di)))[-4:]
stu = ("    " + str(non_pkgs))[-4:]
tpa = ("    " + str(tot_pkgs))[-4:]
tfa = ("    " + str(dl))[-4:]
tpd = ("    " + str(dup_pkgs))[-4:]
tfd = ("    " + str(dup_files))[-4:]

#-----------------------------------------------------------------------------
# Display archive statistics
#-----------------------------------------------------------------------------

print "Total installed packages:   %s   Total uninstalled packages in archive: %s" % (sti, stu)
print "Total packages in archive:  %s   Total files in archive:                %s" % (tpa, tfa)
print "Packages in archive w/dups: %s   Files in archive in dup pkgs:          %s" % (tpd, tfd)

yn = "N\n"

if ((non_pkgs == 0) and (dup_pkgs == 0)):
  print "\n\nNothing to clean in archives.  Exiting.\n\n"
  sys.exit(0)
  
if (non_pkgs > 0):
  print "\n\nDelete uninstalled packages from archive (y/N/p/q)? ",

  yn = sys.stdin.readline()
  if yn == "\n":
    yn = "N\n"
  yn = string.upper(yn[:1])
 
if ((yn == "Y") or (yn == "P")):
#-----------------------------------------------------------------------------
# Remove uninstalled package files from the archive if requested
#-----------------------------------------------------------------------------
  for n in range(dl):
    inst = is_installed(dn[n])
    if (inst == 0):
      if (yn == "P"):
        print "\bKeep:", df[n], " (Y/n/q) ? ",

        yn2 = sys.stdin.readline()
        if yn2 == "\n":
          yn2 = "Y\n"
        yn2 = string.upper(yn2[:1])

        if yn2 == "Q":
          sys.exit(0)

        if yn2 == "N":
          print "\bRemoving: " + basepath + df[n]
          os.remove(basepath + df[n])
      else:
        print "\bRemoving: " + basepath + df[n]
        os.remove(basepath + df[n])

if (yn == "Q"):
  sys.exit(0)

for n in range(1,dl):
#-----------------------------------------------------------------------------
# Determine ranges of duplicate package files for display, prompting,
# and possible removal
#-----------------------------------------------------------------------------

  if dn[n] == dn[n-1]:
    if ms == -1:
      ms = n-1
  else:
    if ms != -1:
      me = n-1

  if (ms > -1) and (n == dl-1):
    me = dl-1

  if (me > -1):
    yn = remove_dups()
    if not yn:
      prompt_dups()

    ms = -1
    me = -1