!7

sclass IPInfo {
  S ip;
  Timestamp firstSeen, lastSeen, blockUntil;
  long passwordFails, passwordSuccesses;
  Timestamp lastPasswordSuccess, lastPasswordFail;
  long blockedRequests;
  Bool allowed;
  
  bool isEvil() { ret blockUntil != null; }
}

static MapSO _renameClasses = litmap("SalterService" := "ScreenshotService");

cm ScreenshotService > DynPrintLogAndEnabled {
  !include early #1029545 // API for Eleu
  
  switchable S password = aGlobalID();
  switchable int badPWBlockSeconds = 30;

  Map<S, IPInfo> ipsSeen = syncLinkedHashMap();
  
  start {  // to show warnings early
    canServe();
    printIPsSeen();
  }
  
  void printIPsSeen aka printStats() {
    var list = cloneValues(ipsSeen);
    print(n2(list, "IP") + " seen." + stringIf(empty(list), " I am a virgin."));
    for (IPInfo info : list)
      print("Known client IP: " + info.ip
        + appendBracketed_comma(
          info.isEvil() ? "evil" : "good",
          info.passwordFails == 0 ? ""
            : n2(info.passwordSuccesses, "normal request") + ", last: "
              + info.lastPasswordSuccess,
          info.passwordFails == 0 ? ""
            : n2(info.passwordFails, "password failures") + ", last: " + info.lastPasswordFail,
          n2UnlessZero(info.blockedRequests, "blocked requests")));
  }
  
  O html(IWebRequest req) ctex {
    if (!canServe()) ret "Can't serve";

    S ip = req.clientIP();
    IPInfo ipInfo = getOrCreate(ipsSeen, ip, () -> {
      print("New client IP!! " + ip);
      new IPInfo info;
      info.ip = ip;
      info.lastSeen = new Timestamp;
      if (info.firstSeen == null) info.firstSeen = info.lastSeen;
      change();
      ret info;
    });
    
    // Don't allow brute-force password checking
    
    if (ipInfo.blockUntil != null && cmp(ipInfo.blockUntil, new Timestamp) > 0) {
      print("Request from blocked IP: " + ip);
      ipInfo.blockedRequests++;
      change();
      ret print("Not allowed (evil IP)");
    }

    S uri = req.uri();
    print("Serving to IP " + ip + ": " + uri);
    
    if (isFalse(ipInfo.allowed))
      ret print("Not allowed (bad IP)");
      
    S suppliedPW = req.get("password");
    if (empty(suppliedPW)) ret "Need password";
    
    if (!eq(password, suppliedPW)) {
      ++ipInfo.passwordFails;
      ipInfo.lastPasswordFail = new Timestamp;
      change();
      print("Password FAIL from " + ip + " OK");
      if ((ipInfo.passwordFails % 10) == 0) {
        infoMessage(n2(ipInfo.passwordFails) + " password fails from " + ip + "!!!!");
        print("Blocking " + ip + " for " + badPWBlockSeconds + " seconds");
        ipInfo.blockUntil = new Timestamp(nowPlusSeconds(badPWBlockSeconds));
        change();
      }
      ret "Bad PW";
    }
    
    ++ipInfo.passwordSuccesses;
    ipInfo.lastPasswordSuccess = new Timestamp;
    change();
    print("Password from " + ip + " OK");

    ret subBot_serveJPEG(shootScreen2());
  }
  
  bool canServe() {
    if (!enabled) ret false with print("Disabled");
    if (l(password) < 4) ret false with warn("Not starting - password too short!!");
    true;
  }
  
  void forgetIPsSeen() {
    ipsSeen.clear();
    change();
    print("Mandatory memory wipe done.");
    printIPsSeen();
  }
  
  enhanceFrame {
    addInternalFrameTitleMenuItem(f/JInternalFrame, "Print Stats", rThreadEnter printStats);
  }
}